dataset.py

import os
import random
import pandas as pd
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset

   
def collate_fn(examples):
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    labels = torch.tensor([example["label"] for example in examples])
    return {"pixel_values": pixel_values, "labels": labels}


class CustomImageDataset(Dataset):
    def __init__(self, filepaths, labels, transform):
        self.filepaths = filepaths
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx):
        imagem = Image.open(self.filepaths[idx]).convert("RGB")
        if self.transform:
            imagem = self.transform(imagem)
        return {
            "pixel_values": imagem,
            "label": int(self.labels[idx])
        }

class CustomTensorDataset(Dataset):
    def __init__(self, caminho_csv, pasta_tensores):
        self.data = pd.read_csv(caminho_csv)
        self.pasta_tensores = pasta_tensores

        # Transform labels to 0 for NC ou 1 for G3, G4 and G5
        if 'ground truth' in self.data.columns:
            labels = self.data['ground truth'].values
        else:
            labels = self.data['MV'].values
        self.labels = (labels != 0).astype(int)  

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

    def __getitem__(self, idx):
        nome_arquivo = self.data.iloc[idx]['Patch filename'] + ".pt"
        tensor_path = os.path.join(self.pasta_tensores, nome_arquivo)

        # LOAD the tensor already processed
        imagem_tensor = torch.load(tensor_path)

        label = int(self.labels[idx])

        return {"pixel_values": imagem_tensor, "label": label}