File size: 1,063 Bytes
35ec68e | 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 | """MS1M dataset: identity folders of 112x112 ArcFace-aligned faces."""
import os
import glob
import cv2
import torch
from torch.utils.data import Dataset
def load_rgb01(path):
img = cv2.imread(path) # BGR uint8, 112x112
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
t = torch.from_numpy(img).float().permute(2, 0, 1) / 255.0
return t
class MS1MDataset(Dataset):
"""Aligned MS1M faces. Uses `num_ids` identity folders (sorted
numerically) starting at `start` so different scripts see the same
deterministic split."""
def __init__(self, root, num_ids, start=0):
ids = sorted(os.listdir(root), key=int)[start:start + num_ids]
self.samples = []
for label, ident in enumerate(ids):
for f in sorted(glob.glob(f"{root}/{ident}/*.jpg")):
self.samples.append((f, label))
self.num_classes = len(ids)
def __len__(self):
return len(self.samples)
def __getitem__(self, i):
path, label = self.samples[i]
return load_rgb01(path), label
|