preprocessing.py

import pandas as pd
import numpy as np
from typing import List, Tuple
from sklearn.preprocessing import LabelEncoder
import torch
from torch.utils.data import TensorDataset
from PIL import Image
from torchvision import transforms
from sklearn.decomposition import PCA

COLS_TO_DROP = ["No.", "Id"]
COL_ID = ["Id"]
COL_LABEL = ["wheatvariety"]
COLS_TO_KEEP = [
    "kernelarea",
    "kernelperimeter",
    "compactness",
    "kernellength",
    "kernelwidth",
    "asymmetry",
    "groovelength",
    "germarea",
    "germlength",
]
COLS_TO_ADD = [
    "germarea_kernelarea",
    "germlength_kernellength",
    "kernelwidth_kernellength",
]
COL_FEATURES = COLS_TO_ADD + COLS_TO_KEEP


def drop_columns(df: pd.DataFrame) -> pd.DataFrame:
    return df.drop(columns=COLS_TO_DROP)


def encode_nominal_data(df: pd.DataFrame) -> pd.DataFrame:
    encoder = LabelEncoder()
    for col in COL_LABEL:
        df[col] = encoder.fit_transform(df[col])
    return df


def normalize_data(df: pd.DataFrame) -> pd.DataFrame:
    for col in COL_FEATURES:
        df[col] = (df[col] - df[col].min()) / (df[col].max() - df[col].min())
    return df


def add_indirect_features(df: pd.DataFrame) -> pd.DataFrame:
    df_new = df.copy()
    df_new["germarea_kernelarea"] = df["germarea"] / df["kernelarea"]
    df_new["germlength_kernellength"] = df["germlength"] / df["kernellength"]
    df_new["kernelwidth_kernellength"] = df["kernelwidth"] / df["kernellength"]
    return df_new


def preprocess_data(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
    df = drop_columns(df)
    df = encode_nominal_data(df)
    df = add_indirect_features(df)
    df = normalize_data(df)

    X = df[COL_FEATURES]
    y = df[COL_LABEL]
    return X, y


def pd_to_numpy_X_y(X: pd.DataFrame, y: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]:
    X_np = X.to_numpy()
    y_np = y.to_numpy().ravel()
    return X_np, y_np


def numpy_to_tensor_X_y(
    X: np.ndarray, y: np.ndarray
) -> Tuple[torch.Tensor, torch.Tensor]:
    X_tensor = torch.tensor(X, dtype=torch.float32)
    y_tensor = torch.tensor(y, dtype=torch.long)
    return X_tensor, y_tensor


def get_tensor_dataset(X: np.ndarray, y: np.ndarray) -> TensorDataset:
    X_tensor, y_tensor = numpy_to_tensor_X_y(X, y)
    return TensorDataset(X_tensor, y_tensor)


### IMAGE PREPROCESSING ########################################


def to_grayscale(img: Image.Image) -> Image.Image:
    return img.convert("L")


def center_crop(img: Image.Image, target_width: int, target_height: int) -> Image.Image:
    width, height = img.size

    left = (width - target_width) // 2
    top = (height - target_height) // 2
    right = left + target_width
    bottom = top + target_height

    return img.crop((left, top, right, bottom))


def pad_image(img: Image.Image, target_width: int, target_height: int) -> Image.Image:
    org_width, org_height = img.size
    if target_height < org_height or target_width < org_width:
        raise ValueError(
            f"Target dimensions ({target_width}, {target_height}) can't be smaller than original dimensions ({org_width}, {org_height})"
        )
    new_image = Image.new(img.mode, (target_width, target_height))
    pad_size_left = (target_width - org_width) // 2
    pad_size_top = (target_height - org_height) // 2
    new_image.paste(img, (pad_size_left, pad_size_top))
    return new_image


def image_to_tensor(img: Image.Image) -> torch.Tensor:
    to_tensor = transforms.ToTensor()
    tensor = to_tensor(img)
    return tensor


def get_max_sizes(images: List[Image.Image]) -> Tuple[int, int]:
    w, h = [], []
    for img in images:
        w.append(img.width)
        h.append(img.height)
    return max(w), max(h)


def get_min_sizes(images: List[Image.Image]) -> Tuple[int, int]:
    w, h = [], []
    for img in images:
        w.append(img.width)
        h.append(img.height)
    return min(w), min(h)


def preprocess_images(images: List[Image.Image]) -> torch.Tensor:
    max_width, max_height = get_max_sizes(images)

    preprocessed_images = []
    for img in images:
        img = to_grayscale(img)
        img = pad_image(img, max_width, max_height)
        img = image_to_tensor(img)
        preprocessed_images.append(img)

    return torch.stack(preprocessed_images, dim=0)


### IMAGE REDUCTION ########################################


def img_tensor_to_ndarray(images: torch.Tensor) -> np.ndarray:
    return images.cpu().detach().numpy()


def dim_reduction(images: np.ndarray | torch.Tensor, target_dim: int) -> torch.Tensor:
    images_flat = images.reshape(len(images), -1)
    pca = PCA(n_components=target_dim)
    images_reduced = pca.fit_transform(images_flat)  # type: ignore
    return torch.from_numpy(images_reduced)


### JOINED PREPROCESSING ########################################


def sort_dataframe(df: pd.DataFrame, ids: List[str]) -> pd.DataFrame:
    ids = [iid[:-4] for iid in ids]
    return df.set_index(COL_ID).loc[ids].reset_index()


def preprocess_all(
    df: pd.DataFrame, imgs: List[Image.Image], ids: List[str]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    df = sort_dataframe(df, ids)
    imgs_tensor = preprocess_images(imgs)
    df_x, df_y = preprocess_data(df)

    numpy_x, numpy_y = pd_to_numpy_X_y(df_x, df_y)
    tensor_x, tensor_y = numpy_to_tensor_X_y(numpy_x, numpy_y)

    return tensor_x, imgs_tensor, tensor_y


def tensors_to_dataset(
    features: torch.Tensor, imgs: torch.Tensor, labels: torch.Tensor
) -> TensorDataset:
    return TensorDataset(features, imgs, labels)


def join_multisource(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
    return torch.cat((x1, x2), dim=1)