src/models/skinLesionDatasets.py

from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import OneHotEncoder, LabelEncoder, StandardScaler
import pandas as pd
import numpy as np
from PIL import Image
from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2
import torch
import os
import pickle
import cv2

class SkinLesionDataset(Dataset):
    def __init__(self, metadata_file:str, img_dir:str, bert_model_name="one-hot-encoder", size:tuple=(224,224),
        drop_nan:bool=False, random_undersampling:bool=False,
        image_encoder:str="resnet-50", is_train:bool=True):
        # Store parameters
        self.metadata_file = metadata_file
        self.img_dir = img_dir
        self.size = size
        self.bert_model_name=bert_model_name
        self.is_to_drop_nan = drop_nan
        self.random_undersampling = random_undersampling
        self.image_encoder = image_encoder
        self.is_train = is_train
        self.targets = None
        self.normalization = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        self.transform = self.load_transforms()

        # Load metadata and process
        self.metadata = self.load_metadata()
        self.features, self.labels, self.targets = self.one_hot_encoding()


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

    def __getitem__(self, idx):
        image_name = self.metadata.iloc[idx]['img_id']
        img_path = os.path.abspath(os.path.join(self.img_dir, image_name))

        try:
            with Image.open(img_path) as img:
                image = img.convert("RGB")
                image = np.array(image)
        except Exception as e:
            print(f"[Erro] Não foi possível abrir imagem com PIL: {img_path} — {e}")
            raise FileNotFoundError(f"Imagem inválida: {img_path}")

        if self.transform:
            image = self.transform(image=image)['image']

        metadata = torch.tensor(self.features[idx], dtype=torch.float32)
        label = torch.tensor(self.labels[idx], dtype=torch.long)

        return image_name, image, metadata, label

    def load_transforms(self):
        """
        Define as transformações de imagem para treino/validação.

        - Treino:
            * Resize fixo
            * Rotate moderado (±45°) – implementação própria do Albumentations (cv2), sem skimage.AffineTransform
            * Flips horizontal/vertical
            * Blur, dropout e variação de cor
            * Normalização + ToTensorV2

        - Val/Test:
            * Apenas Resize + Normalize + ToTensorV2
        """
        if self.is_train:
            return A.Compose([
                # Ajuste de tamanho base
                A.Resize(self.size[0], self.size[1]),

                # Geométricas SEGURAS (sem Affine / ShiftScaleRotate)
                A.Rotate(
                    limit=45,
                    border_mode=cv2.BORDER_REFLECT,
                    p=0.5
                ),

                # Flips
                A.HorizontalFlip(p=0.5),
                A.VerticalFlip(p=0.2),

                # Blur
                A.GaussianBlur(sigma_limit=(0, 2.0), p=0.25),

                # Dropout leve (oclusões pequenas)
                A.CoarseDropout(
                    max_holes=5,
                    max_height=8,
                    max_width=8,
                    p=0.15
                ),

                # Variações de cor/iluminação
                A.HueSaturationValue(
                    hue_shift_limit=10,
                    sat_shift_limit=15,
                    val_shift_limit=10,
                    p=0.25
                ),
                A.RandomBrightnessContrast(p=0.25),

                # Normalização + tensor
                A.Normalize(mean=self.normalization[0], std=self.normalization[1]),
                ToTensorV2(),
            ])
        else:
            # Validação / teste: sem augmentations fortes
            return A.Compose([
                A.Resize(self.size[0], self.size[1]),
                A.Normalize(mean=self.normalization[0], std=self.normalization[1]),
                ToTensorV2(),
            ])


    def load_metadata(self):
        # Carregar o CSV
        metadata = pd.read_csv(self.metadata_file).fillna("EMPTY").replace(" ", "EMPTY").replace("  ", "EMPTY").\
           replace("NÃO  ENCONTRADO", "EMPTY").replace("BRASIL","BRAZIL")
        # Verificar se deve descartar linhas com NaN
        if self.is_to_drop_nan:
            metadata = metadata.dropna().reset_index(drop=True)

        return metadata

    def one_hot_encoding(self):
        dataset_features = self.metadata.drop(
            columns=['patient_id', 'lesion_id', 'img_id', 'biopsed', 'diagnostic']
        )

        # Colunas numéricas fixas
        numerical_cols = ['age', 'diameter_1', 'diameter_2']
        categorical_cols = [col for col in dataset_features.columns if col not in numerical_cols]

        # Converter categóricas
        dataset_features[categorical_cols] = dataset_features[categorical_cols].astype(str)

        # Forçar numérico nas colunas numéricas, substituindo inválidos por NaN
        dataset_features[numerical_cols] = dataset_features[numerical_cols].apply(
            pd.to_numeric, errors="coerce"
        )

        # Preencher valores faltantes (NaN gerados acima) com -1
        dataset_features[numerical_cols] = dataset_features[numerical_cols].fillna(-1)


        # Caminho base
        base_dir = os.path.join("./data", "preprocess_data")
        os.makedirs(base_dir, exist_ok=True)

        # OneHotEncoder
        ohe_path = os.path.join(base_dir, "ohe_pad_20.pickle")
        if os.path.exists(ohe_path):
            with open(ohe_path, "rb") as f:
                ohe = pickle.load(f)
            categorical_data = ohe.transform(dataset_features[categorical_cols])
        else:
            ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
            categorical_data = ohe.fit_transform(dataset_features[categorical_cols])
            with open(ohe_path, "wb") as f:
                pickle.dump(ohe, f)

        # StandardScaler
        scaler_path = os.path.join(base_dir, "scaler_pad_20.pickle")
        if os.path.exists(scaler_path):
            with open(scaler_path, "rb") as f:
                scaler = pickle.load(f)
            numerical_data = scaler.transform(dataset_features[numerical_cols])
        else:
            scaler = StandardScaler()
            numerical_data = scaler.fit_transform(dataset_features[numerical_cols])
            with open(scaler_path, "wb") as f:
                pickle.dump(scaler, f)

        # Concatenar dados
        processed_data = np.hstack((categorical_data, numerical_data))

        # Labels
        labels = self.metadata['diagnostic'].values
        le_path = os.path.join(base_dir, "label_encoder_pad_20.pickle")
        if os.path.exists(le_path):
            with open(le_path, "rb") as f:
                label_encoder = pickle.load(f)
            encoded_labels = label_encoder.transform(labels)
        else:
            label_encoder = LabelEncoder()
            encoded_labels = label_encoder.fit_transform(labels)
            with open(le_path, "wb") as f:
                pickle.dump(label_encoder, f)

        return processed_data, encoded_labels, self.metadata['diagnostic'].unique()