AestheticPredictor/pytorch/data.py

import os
import torch
import torchvision.models as models

from PIL import Image
from torch.utils.data import Dataset


class EfficientNetFeatureExtractor:
    def __init__(self):
        model = models.efficientnet_v2_m(weights=models.EfficientNet_V2_M_Weights.IMAGENET1K_V1)
        model = torch.nn.Sequential(*list(model.children())[:-1])
        
        self.model = model
        self.preprocess = models.EfficientNet_V2_M_Weights.IMAGENET1K_V1.transforms()
        self.feature_dims = 1280

    def extract(self, image_path):
        with torch.no_grad():
            image = Image.open(image_path).convert('RGB')
            X = self.preprocess(image)
            X = X.unsqueeze(0)
            embedding = self.model(X)[0, :, 0, 0]
        return embedding


class LAIONAestheticsDataset(Dataset):
    def __init__(self, annotations_file, img_dir, feature_extractor: EfficientNetFeatureExtractor, limit=None):
        
        self.feature_extractor = feature_extractor
        self.img_path = img_dir
        self.img_files = []
        self.scores = []
        self.embeddings = {}

        with open(annotations_file) as f:
            for i, row in enumerate(f.readlines()):
                if limit is not None and i >= limit:
                    break
                img_name, _, aesthetic_score = row.split('\t')[:3]
                self.img_files.append(img_name)
                self.scores.append(torch.tensor([float(aesthetic_score)]))

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        score = self.scores[idx]
        embedding = self.embeddings.get(idx, None)
        if embedding is None:
            img_path = os.path.join(self.img_path, self.img_files[idx])
            embedding = self.feature_extractor.extract(img_path)
            self.embeddings[idx] = embedding

        return embedding, score


def train_valid_split(data_dir, train_percent=0.8, limit=None):
    feature_extractor = EfficientNetFeatureExtractor()
    annotations_file = os.path.join(data_dir, 'labels.tsv')
    dataset_train_valid = LAIONAestheticsDataset(annotations_file, data_dir, feature_extractor, limit=limit)
    
    # Split into train/valid
    train_size = int(train_percent * len(dataset_train_valid))
    valid_size = len(dataset_train_valid) - train_size
    train_dataset, valid_dataset = torch.utils.data.random_split(dataset_train_valid, [train_size, valid_size])

    return train_dataset, valid_dataset