| 107 | |
| 108 | |
| 109 | class RandomMultipleGallerySamplerNoCam(Sampler): |
| 110 | def __init__(self, data_source, num_instances=4): |
| 111 | super().__init__(data_source) |
| 112 | |
| 113 | self.data_source = data_source |
| 114 | self.index_pid = defaultdict(int) |
| 115 | self.pid_index = defaultdict(list) |
| 116 | self.num_instances = num_instances |
| 117 | |
| 118 | for index, _item in enumerate(data_source): |
| 119 | pid = _item[1] |
| 120 | # (_, pid, cam) |
| 121 | if pid < 0: |
| 122 | continue |
| 123 | self.index_pid[index] = pid |
| 124 | self.pid_index[pid].append(index) |
| 125 | |
| 126 | self.pids = list(self.pid_index.keys()) |
| 127 | self.num_samples = len(self.pids) |
| 128 | |
| 129 | def __len__(self): |
| 130 | return self.num_samples * self.num_instances |
| 131 | |
| 132 | def __iter__(self): |
| 133 | indices = torch.randperm(len(self.pids)).tolist() |
| 134 | ret = [] |
| 135 | |
| 136 | for kid in indices: |
| 137 | i = random.choice(self.pid_index[self.pids[kid]]) |
| 138 | # _, i_pid, i_cam = self.data_source[i] |
| 139 | # _, i_pid, i_cam |
| 140 | ret.append(i) |
| 141 | |
| 142 | pid_i = self.index_pid[i] |
| 143 | index = self.pid_index[pid_i] |
| 144 | |
| 145 | select_indexes = No_index(index, i) |
| 146 | if not select_indexes: |
| 147 | continue |
| 148 | if len(select_indexes) >= self.num_instances: |
| 149 | ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=False) |
| 150 | else: |
| 151 | ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=True) |
| 152 | |
| 153 | for kk in ind_indexes: |
| 154 | ret.append(index[kk]) |
| 155 | |
| 156 | return iter(ret) |