blob: 1dd33ca5f0d14d8d442592031144e4ae3d1e80c2 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
import random
from typing import Iterator, Tuple
import numpy as np
from torch.utils import data
from utils.dataset import CASIAB
class TripletSampler(data.Sampler):
def __init__(
self,
data_source: CASIAB,
batch_size: Tuple[int, int]
):
super().__init__(data_source)
self.metadata_label = data_source.metadata['labels']
self.labels = data_source.labels
self.length = len(self.labels)
self.indexes = np.arange(0, self.length)
(self.P, self.K) = batch_size
def __iter__(self) -> Iterator[int]:
while True:
sampled_indexes = []
sampled_labels = random.sample(self.metadata_label, k=self.P)
for label in sampled_labels:
clip_indexes = list(self.indexes[self.labels == label])
# Sample without replacement if have enough clips
if len(clip_indexes) >= self.K:
_sampled_indexes = random.sample(clip_indexes, k=self.K)
else:
_sampled_indexes = random.choices(clip_indexes, k=self.K)
sampled_indexes += _sampled_indexes
yield sampled_indexes
def __len__(self) -> int:
return self.length
|
