ModelThatExtractsOverallArtCharacteristics.py

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import pandas as pd
import os
from pathlib import Path

# ---------------------- CONFIG ----------------------
BATCH_SIZE = 32
EPOCHS = 100
LR = 1e-4
LABEL_COLUMNS = [
    'Total aesthetic score', 'Theme and logic', 'Creativity',
    'Layout and composition', 'Space and perspective', 'The sense of order',
    'Light and shadow', 'Color', 'Details and texture', 'The overall', 'Mood'
]
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -------------------- Dataset --------------------
class APDDv2Dataset(Dataset):
    def __init__(self, csv_path, image_dir, transform=None):
        self.data = pd.read_csv(csv_path)
        self.image_dir = image_dir
        self.transform = transform
        self.label_columns = LABEL_COLUMNS

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        img_path = os.path.join(self.image_dir, row['filename'])

        if not os.path.exists(img_path):
            return None

        try:
            img = Image.open(img_path).convert("RGB")
            if self.transform:
                img = self.transform(img)

            labels = []
            mask = []
            for col in self.label_columns:
                val = row.get(col, None)
                if pd.isna(val) or val == '':
                    labels.append(0.0)
                    mask.append(0.0)
                else:
                    labels.append(float(val) / 10.0)  # Normalize to 0–1
                    mask.append(1.0)

            return img, torch.tensor(labels, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32)
        except:
            return None

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

    def collate_fn(self, batch):
        batch = [b for b in batch if b is not None]
        if len(batch) == 0:
            return None
        images = torch.stack([b[0] for b in batch])
        labels = torch.stack([b[1] for b in batch])
        masks = torch.stack([b[2] for b in batch])
        return images, labels, masks

# -------------------- Transforms --------------------
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# -------------------- Dataset and Loader --------------------
csv_path = "FarmDataset.csv"
image_dir = "cached_images"
dataset = APDDv2Dataset(csv_path, image_dir, transform=train_transform)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=dataset.collate_fn)

data_path = Path("CrystalDataset/")

train_dir = data_path / "train"
test_dir = data_path / "test"

# -------------------- Model --------------------
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
for param in model.parameters():
    param.requires_grad = False
model.fc = nn.Sequential(
    nn.Linear(2048, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, len(LABEL_COLUMNS))
)
model.to(DEVICE)

# -------------------- Optimizer & Loss --------------------
optimizer = torch.optim.Adam(model.fc.parameters(), lr=LR)
criterion = nn.SmoothL1Loss(reduction='none')  # We'll apply mask manually
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=True)

# -------------------- Training Loop --------------------
for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0
    batches = 0

    for batch in dataloader:
        if batch is None:
            continue
        images, labels, masks = batch
        images, labels, masks = images.to(DEVICE), labels.to(DEVICE), masks.to(DEVICE)

        optimizer.zero_grad()
        outputs = model(images)
        loss_raw = criterion(outputs, labels)
        loss = (loss_raw * masks).sum() / masks.sum()
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()
        batches += 1

    if batches > 0:
        avg_loss = epoch_loss / batches
        print(f"Epoch {epoch+1}/{EPOCHS} - Loss: {avg_loss:.4f}")
        scheduler.step(avg_loss)
    else:
        print(f"Epoch {epoch+1}: No valid data processed.")


model.eval()
results = []
with torch.inference_mode():
    folders_path_list = [p for p in train_dir.iterdir()]
    for folder in folders_path_list:
        for image_path in os.listdir(folder):
            img = Image.open(folder.joinpath(image_path)).convert("RGB")
            img = train_transform(img).unsqueeze(0).to(DEVICE)  # Already here, good!
            outputs = model(img)
            outputs = outputs.cpu().numpy()[0]
            results.append([folder.joinpath(image_path)] + outputs.tolist())

    folders_path_list_test = [p for p in test_dir.iterdir()]
    for folder in folders_path_list_test:
        for image_path in os.listdir(folder):
            img = Image.open(folder.joinpath(image_path)).convert("RGB")
            img = train_transform(img).unsqueeze(0).to(DEVICE)  # Already here, good!
            outputs = model(img)
            outputs = outputs.cpu().numpy()[0]
            results.append([folder.joinpath(image_path)] + outputs.tolist())

results_df = pd.DataFrame(
    results,
    columns=[
        "image",
        "Total aesthetic score",
        "Theme and logic",
        "Creativity",
        "Layout and composition",
        "Space and perspective",
        "Disgust",
        "Light and shadow",
        "Color",
        "Details and texture",
        "The overall",
        "Mood",
    ],
)
results_df.to_csv("FarmedAnimals.csv", index=False)