| # Class used in extract_features.py file | |
| from torch.utils.data import Dataset | |
| from PIL import Image | |
| import torch | |
| from torchvision import transforms | |
| import os | |
| import cv2 | |
| class FacesDataset(Dataset): | |
| def __init__(self, faces, output_dir) -> None: | |
| super().__init__() | |
| self.faces = faces | |
| self.output_dir = output_dir | |
| def __getitem__(self, index: int): | |
| # Preprocess the image as required by EfficientNet | |
| face_path = self.faces[index] | |
| tfms = transforms.Compose([transforms.Resize((224,224)), transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),]) | |
| img = tfms(Image.open(face_path)) | |
| # Compose the output path | |
| output_path = self.output_dir + face_path.split("faces")[1] + ".pt" | |
| return img, output_path | |
| def __len__(self) -> int: | |
| return len(self.faces) | |