glisicstefan/age-estimation-resnet50

1import torch2import os3from torch.utils.data import Dataset, DataLoader, random_split4from PIL import Image5 6class UTKFaceDataset(Dataset):7    def __init__(self, root_dir, transform=None):8        self.root_dir = root_dir9        self.transform = transform10 11        self.img_names = [f for f in os.listdir(root_dir) if (10 < int(f.split("_")[0]) < 65)]12 13    def __len__(self):14        return len(self.img_names)15    16    def __getitem__(self, idx):17        img_name = self.img_names[idx]18        img_path = os.path.join(self.root_dir, img_name)19 20        image = Image.open(img_path).convert("RGB")21        age = int(img_name.split("_")[0])22 23        if self.transform:24            image = self.transform(image)25 26        return image, torch.tensor(age, dtype=torch.float32)27    28class ApplyTransform(Dataset):29    """30    Omotava dataset (ili subset) i primenjuje specifičnu transformaciju.31    """32    def __init__(self, subset, transform=None):33        self.subset = subset34        self.transform = transform35        36    def __getitem__(self, index):37        x, y = self.subset[index]38        if self.transform:39            x = self.transform(x)40        return x, y41        42    def __len__(self):43        return len(self.subset)44    45 46def create_dataloaders(root_dir, train_transform, test_transform, batch_size=32, train_split=0.8):47  48    full_dataset = UTKFaceDataset(root_dir=root_dir, transform=None)49    50 51    train_len = int(len(full_dataset) * train_split)52    test_len = len(full_dataset) - train_len53    train_subset, test_subset = random_split(full_dataset, [train_len, test_len])54    55    train_data = ApplyTransform(train_subset, transform=train_transform)56    test_data = ApplyTransform(test_subset, transform=test_transform)57    58    train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True)59    test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)60    61    return train_dataloader, test_dataloader62