File size: 897 Bytes
9e14838 | 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 | # 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)
|