| """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 | |