train.py

import csv
import os
import sys
import shutil

import pandas as pd
import numpy as np
from math import isnan, inf

import cv2
import copy

from scipy.stats import skew
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

from config import LEARNING_RATE, EPOCHS, BATCH_SIZE, PATIENCE
from Augmentation import AUGMENTS
from Network import Network
from Transformation import BatchTransformator

SAVE_DIR = "dataset"


def get_images(directory):
    images = []
    files = os.listdir(directory)
    for file in files:
        file_path = os.path.join(directory, file)
        if os.path.isfile(file_path) and file_path.lower().endswith("jpg"):
            img = cv2.imread(file_path)
            if img is None:
                raise PermissionError
            images.append({"path": file, "image": img})
    return images


def get_augmented_images(base_images):
    augmented_images = []
    for image in base_images:
        image_augments = [image]
        filename, ext = os.path.splitext(os.path.basename(image["path"]))
        for augment_name, augment_fun in AUGMENTS.items():
            image_augments.append(
                {"path": f"{filename}_{augment_name}{ext}",
                 "image": augment_fun(image=image["image"])["image"]})
        augmented_images.extend(image_augments)
    return augmented_images


def get_transformations(base_images):
    print("\n⏳ Transforming...")
    bt = BatchTransformator()

    imgs = [img["image"] for img in base_images]
    paths = [img["path"] for img in base_images]
    bt.import_images(imgs, paths)
    bt.apply_transformations()

    return bt


def get_histogram_features(t):
    labels = ["red", "green", "blue", "saturation", "value",
              "lightness", "green-magenta", "blue-yellow"]
    methods = [("skew", lambda x: skew(x)),
               ("p1", lambda x: np.percentile(x, 50)),
               ("p2", lambda x: np.percentile(x, 75)),
               ("p3", lambda x: np.percentile(x, 90)),
               ("std", np.std)]
    features = {
        f"{label}_{method_name}":
            method(t.hist_data[f"{label}_frequencies"]["value"])
        for method_name, method in methods for label in labels
    }
    features["hue_circular_std"] = t.hist_data["hue_circular_std"]["value"]
    features["hue_median"] = t.hist_data["hue_median"]["value"]
    features["hue_circular_mean"] = t.hist_data["hue_circular_mean"]["value"]

    return features


def get_analyze_features(t):
    features = t.analyze_data
    return features


def add_features(features, new_features):
    for key, value in new_features.items():
        if isinstance(features.get(key), list):
            features[key].append(value)
        else:
            features[key] = [value]


def calculate_features(bt, label):
    features = {"group": []}

    for t in bt.transformations:
        features["group"].append(label)
        add_features(features, get_histogram_features(t))
        add_features(features, get_analyze_features(t))

    return features


def save_dataset(images):
    try:
        os.mkdir(SAVE_DIR)
    except FileExistsError:
        pass

    print("➡️ Saving dataset...")
    for key, bt in images.items():
        print(f"\n⏳ Saving {key}")
        bt.save_images(f"{SAVE_DIR}/{key}")
    shutil.make_archive(SAVE_DIR, "zip", SAVE_DIR)


def prepare_dataset(directory):
    features = {}

    try:
        dir_files = os.listdir(directory)
        for file_path in dir_files:
            path = os.path.join(directory, file_path)
            if os.path.isdir(path):
                print(f"➡️ Now processing {file_path}")

                print("➡️ Reading images...")
                images = get_images(path)

                print("➡️ Augmenting images...")
                augmented_images = get_augmented_images(images)

                print("➡️ Transforming images...")
                bt = get_transformations(augmented_images)

                print("➡️ Calculating features...\n")
                bt_features = calculate_features(bt, file_path)

                for key, value in bt_features.items():
                    if key not in features:
                        features[key] = value
                    else:
                        features[key].extend(value)

    except NotADirectoryError:
        return print("❌ Error: not a directory.")
    except FileNotFoundError:
        return print("❌ Error: directory could not be found.")
    except PermissionError:
        return print("❌ Error: permission denied on directory")

    if "group" not in features:
        return print("❌ Error: no features")

    min_maxes = {}
    xmins = {}

    for key, array in features.items():
        if key == "group":
            continue
        np_arr = np.array(array)

        nan_mask = np.isnan(np_arr)
        if nan_mask.any():
            median_val = np.nanmedian(np_arr)
            np_arr[nan_mask] = median_val if not isnan(median_val) else 0.0
        xmin = np_arr.min()
        xmax = np_arr.max()
        min_max = xmax - xmin
        features[key] = (np_arr - xmin) / min_max if min_max != 0\
            else np_arr - xmin

        min_maxes[key] = min_max
        xmins[key] = xmin

    with open("normalization.csv", "w") as f:
        writer = csv.DictWriter(f, fieldnames=min_maxes.keys())
        writer.writeheader()
        writer.writerows([min_maxes, xmins])

    with open("dataset.csv", "w") as f:
        rows = np.array([{f: features[f][i] for f in features.keys()}
                         for i, _ in enumerate(features["group"])])
        np.random.shuffle(rows)
        writer = csv.DictWriter(f, fieldnames=features.keys())
        writer.writeheader()
        writer.writerows(rows)

    return features


def prepare_dataset_csv(directory):
    df_train = pd.read_csv("train.csv")
    df_val = pd.read_csv("validation.csv")

    y_train = df_train["group"].values
    X_train = df_train.drop(columns=["group"]).values

    y_val = df_val["group"].values
    X_val = df_val.drop(columns=["group"]).values

    X_train = torch.tensor(X_train, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.long)

    X_val = torch.tensor(X_val, dtype=torch.float32)
    y_val = torch.tensor(y_val, dtype=torch.long)

    train_dataset = TensorDataset(X_train, y_train)
    val_dataset = TensorDataset(X_val, y_val)

    input_size = X_train.shape[1]
    num_classes = len(torch.unique(y_train))

    model = Network(input_size, num_classes)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, shuffle=True, batch_size=BATCH_SIZE)

    epoch_count = EPOCHS

    best_val_loss = inf
    curr_patience = 0

    for epoch in range(epoch_count):  # Training for 100 epochs
        running_loss = 0.0  # Track total loss for each epoch
        val_running_loss = 0.0

        # Loop over training batches
        for inputs, labels in loader:
            optimizer.zero_grad()  # Clear previous gradients
            outputs = model(inputs)  # Forward pass
            loss = criterion(outputs, labels)  # Compute loss

            loss.backward()  # Backpropagation
            optimizer.step()  # Update weights

            running_loss += loss.item()  # Accumulate loss

        model.eval()
        with torch.no_grad():
            for val_inputs, val_labels in val_loader:
                val_outputs = model(val_inputs)
                val_loss = criterion(val_outputs, val_labels)
                val_running_loss += val_loss.item()

        model.train()

        # Log loss for the current epoch
        avg_loss = running_loss / len(loader)
        avg_val_loss = val_running_loss / len(val_loader)
        print(f"Epoch [{epoch + 1}/{epoch_count}], Loss: {avg_loss:.4f}, "
              + f"Val loss: {avg_val_loss:.4f}, Best: {best_val_loss:.4f}")

        avg_val_loss = round(avg_val_loss, 4)
        if best_val_loss > avg_val_loss:
            best_val_loss = avg_val_loss
            best_model = copy.deepcopy(model.state_dict())
            curr_patience = 0
        else:
            curr_patience += 1
            if curr_patience > PATIENCE:
                break

    torch.save(best_model, "model")

    return


def train(features):
    _, arr = np.unique(features["group"], return_counts=True)
    f_count = len(arr)

    model = Network(len(features.keys()) - 1, f_count)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    groups = features["group"]
    del features["group"]

    unique, encoded = np.unique(groups, return_inverse=True)
    groups = encoded

    X = np.column_stack(list(features.values()))
    X = torch.tensor(X, dtype=torch.float32)
    y = torch.tensor(groups)

    dataset = TensorDataset(X, y)

    train, validation = torch.utils.data.random_split(dataset, [0.7, 0.3])

    X_val = torch.stack([validation[i][0] for i in range(len(validation))])
    y_val = torch.tensor([validation[i][1] for i in range(len(validation))])

    df = pd.DataFrame(X_val.numpy())
    df["group"] = y_val.numpy()

    X_train = torch.stack([train[i][0] for i in range(len(train))])
    y_train = torch.tensor([train[i][1] for i in range(len(train))])

    df_train = pd.DataFrame(X_train.numpy())
    df_train["group"] = y_train.numpy()

    df.to_csv("validation.csv", index=False)
    df_train.to_csv("train.csv", index=False)

    loader = DataLoader(train, shuffle=True, batch_size=BATCH_SIZE)
    val_loader = DataLoader(validation, shuffle=True, batch_size=BATCH_SIZE)

    epoch_count = EPOCHS
    best_val_loss = inf
    curr_patience = 0

    for epoch in range(epoch_count):
        running_loss = 0.0
        val_running_loss = 0.0

        for inputs, labels in loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        model.eval()

        with torch.no_grad():
            for val_inputs, val_labels in val_loader:
                val_outputs = model(val_inputs)
                val_loss = criterion(val_outputs, val_labels)
                val_running_loss += val_loss.item()

        model.train()

        avg_loss = running_loss / len(loader)
        avg_val_loss = val_running_loss / len(val_loader)
        print(f"Epoch [{epoch + 1}/{epoch_count}], Loss: {avg_loss:.4f}, "
              f"Val loss: {avg_val_loss:.4f}")

        avg_val_loss = round(avg_val_loss, 4)
        if best_val_loss > avg_val_loss:
            best_val_loss = avg_val_loss
            curr_patience = 0
        else:
            curr_patience += 1
            if curr_patience > PATIENCE:
                break

    torch.save(model.state_dict(), "model")

    return


def main():
    if len(sys.argv) != 2:
        return print("usage: train.py <directory>")
    directory = sys.argv[1]

    features = prepare_dataset(directory)
    # features = prepare_dataset_csv(directory)
    if features is None:
        return
    train(features)

    return


if __name__ == '__main__':
    main()