Sample person identities evenly in each batch. Args: train_color_label, train_thermal_label: labels of two modalities color_pos, thermal_pos: positions of each identity batchSize: batch size
| 61 | |
| 62 | |
| 63 | class IdentitySampler(Sampler): |
| 64 | """Sample person identities evenly in each batch. |
| 65 | Args: |
| 66 | train_color_label, train_thermal_label: labels of two modalities |
| 67 | color_pos, thermal_pos: positions of each identity |
| 68 | batchSize: batch size |
| 69 | """ |
| 70 | |
| 71 | def __init__(self, train_color_label, train_thermal_label, color_pos, thermal_pos, num_pos, batchSize, epoch): |
| 72 | uni_label = np.unique(train_color_label) |
| 73 | self.n_classes = len(uni_label) |
| 74 | |
| 75 | N = np.maximum(len(train_color_label), len(train_thermal_label)) |
| 76 | for j in range(int(N / (batchSize * num_pos)) + 1): |
| 77 | batch_idx = np.random.choice(uni_label, batchSize, replace=False) |
| 78 | for i in range(batchSize): |
| 79 | while len(color_pos[batch_idx[i]]) < 4 or len(thermal_pos[batch_idx[i]]) < 4: |
| 80 | batch_idx[i] = np.random.choice(uni_label, 1, replace=False) |
| 81 | print("re-sampling") |
| 82 | sample_color = np.random.choice(color_pos[batch_idx[i]], num_pos) |
| 83 | sample_thermal = np.random.choice(thermal_pos[batch_idx[i]], num_pos) |
| 84 | |
| 85 | if j == 0 and i == 0: |
| 86 | index1 = sample_color |
| 87 | index2 = sample_thermal |
| 88 | else: |
| 89 | index1 = np.hstack((index1, sample_color)) |
| 90 | index2 = np.hstack((index2, sample_thermal)) |
| 91 | self.index1 = index1 |
| 92 | self.index2 = index2 |
| 93 | self.N = N |
| 94 | |
| 95 | def __iter__(self): |
| 96 | return iter(np.arange(len(self.index1))) |
| 97 | |
| 98 | def __len__(self): |
| 99 | return self.N |
| 100 | |
| 101 | |
| 102 | class AverageMeter(object): |