Update

augment.py

@@ -1,50 +0,0 @@
-import os
-import Augmentor
-import shutil
-from configs import *
-
-tasks = ["1", "2", "3", "4", "5", "6"]
-
-for task in tasks:
-    # Loop through all folders in Task 1 and generate augmented images for each class
-    for disease in os.listdir("data/train/raw/Task " + task):
-        if disease != ".DS_Store":
-            print("Augmenting images in class: ", disease, " in Task ", task)
-            # Create a temp folder to combine the raw data and the external data
-            if not os.path.exists(f"data/temp/Task {task}/{disease}/"):
-                os.makedirs(f"data/temp/Task {task}/{disease}/")
-            for file in os.listdir(f"data/train/raw/Task {task}/{disease}"):
-                shutil.copy(f"data/train/raw/Task {task}/{disease}/{file}", f"data/temp/Task {task}/{disease}/{file}")
-            for file in os.listdir(f"data/train/external/Task {task}/{disease}"):
-                shutil.copy(f"data/train/external/Task {task}/{disease}/{file}", f"data/temp/Task {task}/{disease}/{file}")
-            p = Augmentor.Pipeline(f"data/temp/Task {task}/{disease}", output_directory=f"{disease}/", save_format="png")
-            p.rotate(probability=0.8, max_left_rotation=5, max_right_rotation=5)
-            p.flip_left_right(probability=0.8)
-            p.zoom_random(probability=0.8, percentage_area=0.8)
-            p.flip_top_bottom(probability=0.8)
-            p.random_brightness(probability=0.8, min_factor=0.5, max_factor=1.5)
-            p.random_contrast(probability=0.8, min_factor=0.5, max_factor=1.5)
-            p.random_color(probability=0.8, min_factor=0.5, max_factor=1.5)
-            # Generate 100 - total of original images so that the total number of images in each class is 100
-            p.sample(100 - len(p.augmentor_images))
-            # Move the folder to data/train/Task 1/augmented
-            # Create the folder if it does not exist
-            if not os.path.exists(f"data/train/augmented/Task {task}/"):
-                os.makedirs(f"data/train/augmented/Task {task}/")
-            # Move all images in the data/train/Task 1/i folder to data/train/Task 1/augmented/i
-            os.rename(
-                f"data/temp/Task {task}/{disease}/{disease}",
-                f"data/train/augmented/Task {task}/{disease}",
-            )
-            # Rename all the augmented images to [01, 02, 03]
-            number = 0
-            for file in os.listdir(f"data/train/augmented/Task {task}/{disease}"):
-                number = int(number) + 1
-                if len(str(number)) == 1:
-                    number = "0" + str(number)
-                os.rename(
-                    f"data/train/augmented/Task {task}/{disease}/{file}",
-                    f"data/train/augmented/Task {task}/{disease}/{number}.png",
-                )
-
-

configs.py

@@ -1,43 +0,0 @@
-import os
-import torch
-from torchvision import transforms
-from torch.utils.data import Dataset
-from models import *
-
-# Constants
-RANDOM_SEED = 123
-BATCH_SIZE = 32
-NUM_EPOCHS = 100
-LEARNING_RATE = 0.001
-STEP_SIZE = 10
-GAMMA = 0.5
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-NUM_PRINT = 100
-TASK = 1
-RAW_DATA_DIR = r"data/train/raw/Task " + str(TASK)
-AUG_DATA_DIR = r"data/train/augmented/Task " + str(TASK)
-EXTERNAL_DATA_DIR = r"data/train/external/Task " + str(TASK)
-NUM_CLASSES = 7
-MODEL_SAVE_PATH = "output/checkpoints/model.pth"
-MODEL = mobilenet_v3_small(num_classes=NUM_CLASSES)
-
-preprocess = transforms.Compose(
-    [
-        transforms.Resize((64, 64)),  # Resize images to 64x64
-        transforms.ToTensor(),  # Convert to tensor
-        # Normalize 3 channels
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-    ]
-)
-
-# Custom dataset class
-class CustomDataset(Dataset):
-    def __init__(self, dataset):
-        self.data = dataset
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        img, label = self.data[idx]
-        return img, label

data_loader.py

@@ -1,32 +0,0 @@
-from configs import *
-from torchvision.datasets import ImageFolder
-from torch.utils.data import random_split, DataLoader, Dataset
-
-
-def load_data(raw_dir, augmented_dir, external_dir, preprocess):
-    # Load the dataset using ImageFolder
-    raw_dataset = ImageFolder(root=raw_dir, transform=preprocess)
-    external_dataset = ImageFolder(root=external_dir, transform=preprocess)
-    augmented_dataset = ImageFolder(root=augmented_dir, transform=preprocess)
-    dataset = raw_dataset + external_dataset + augmented_dataset
-
-    print("Classes: ", *raw_dataset.classes, sep = ', ')
-    print("Length of raw dataset: ", len(raw_dataset))
-    print("Length of external dataset: ", len(external_dataset))
-    print("Length of augmented dataset: ", len(augmented_dataset))
-    print("Length of total dataset: ", len(dataset))
-
-    # Split the dataset into train and validation sets
-    train_size = int(0.8 * len(dataset))
-    val_size = len(dataset) - train_size
-    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
-
-    # Create data loaders for the custom dataset
-    train_loader = DataLoader(
-        CustomDataset(train_dataset), batch_size=BATCH_SIZE, shuffle=True, num_workers=0
-    )
-    valid_loader = DataLoader(
-        CustomDataset(val_dataset), batch_size=BATCH_SIZE, num_workers=0
-    )
-    
-    return train_loader, valid_loader

eval.py

@@ -1,86 +0,0 @@
-import os
-import torch
-from torchvision.transforms import transforms
-from sklearn.metrics import f1_score
-import pathlib
-from PIL import Image
-from torchmetrics import ConfusionMatrix
-import matplotlib.pyplot as plt
-from configs import *
-from data_loader import load_data  # Import the load_data function
-
-image_path = "data/test/Task 1/"
-
-# Constants
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# Load the model
-MODEL = MODEL.to(DEVICE)
-MODEL.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=DEVICE))
-MODEL.eval()
-
-# Get class labels from the dataset
-class_labels = os.listdir(image_path)
-
-# Define transformation for preprocessing
-preprocess = transforms.Compose(
-    [
-        transforms.Resize((64, 64)),  # Resize images to 64x64
-        transforms.ToTensor(),  # Convert to tensor
-        transforms.Normalize((0.5,), (0.5,)),  # Normalize (for grayscale)
-    ]
-)
-
-def predict_image(image_path, model, transform):
-    model.eval()
-    correct_predictions = 0
-    total_predictions = len(images)
-
-    # Get a list of image files
-    images = list(pathlib.Path(image_path).rglob("*.png"))
-
-    true_classes = []
-    predicted_labels = []
-
-    with torch.no_grad():
-        for image_file in images:
-            print('---------------------------')
-            # Check the true label of the image by checking the sequence of the folder in Task 1
-            true_class = class_labels.index(image_file.parts[-2])
-            print("Image path:", image_file)
-            print("True class:", true_class)
-            image = Image.open(image_file).convert('RGB')
-            image = transform(image).unsqueeze(0)
-            image = image.to(DEVICE)
-            output = model(image)
-            predicted_class = torch.argmax(output, dim=1).item()
-            # Print the predicted class
-            print("Predicted class:", predicted_class)
-            # Append true and predicted labels to their respective lists
-            true_classes.append(true_class)
-            predicted_labels.append(predicted_class)
-
-            # Check if the prediction is correct
-            if predicted_class == true_class:
-                correct_predictions += 1
-
-    # Calculate accuracy and f1 score
-    accuracy = correct_predictions / total_predictions
-    print("Accuracy:", accuracy)
-    f1 = f1_score(true_classes, predicted_labels, average='weighted')
-    print("Weighted F1 Score:", f1)
-
-    # Convert the lists to tensors
-    predicted_labels_tensor = torch.tensor(predicted_labels)
-    true_classes_tensor = torch.tensor(true_classes)
-
-    # Create a confusion matrix
-    conf_matrix = ConfusionMatrix(num_classes=NUM_CLASSES, task='multiclass')
-    conf_matrix.update(predicted_labels_tensor, true_classes_tensor)
-
-    # Plot the confusion matrix
-    conf_matrix.plot()
-    plt.show()
-
-# Call predict_image function
-predict_image(image_path, MODEL, preprocess)

models.py

@@ -1,36 +0,0 @@
-#######################################################
-# This file stores all the models used in the project.#
-#######################################################
-
-# Import all models from torchvision.models
-from torchvision.models import resnet50
-from torchvision.models import resnet18
-from torchvision.models import squeezenet1_0
-from torchvision.models import vgg16
-from torchvision.models import alexnet
-from torchvision.models import densenet121
-from torchvision.models import googlenet
-from torchvision.models import inception_v3
-from torchvision.models import mobilenet_v2
-from torchvision.models import mobilenet_v3_small
-from torchvision.models import mobilenet_v3_large
-from torchvision.models import shufflenet_v2_x0_5
-from torchvision.models import vgg11
-from torchvision.models import vgg11_bn
-from torchvision.models import vgg13
-from torchvision.models import vgg13_bn
-from torchvision.models import vgg16_bn
-from torchvision.models import vgg19_bn
-from torchvision.models import vgg19
-from torchvision.models import wide_resnet50_2
-from torchvision.models import wide_resnet101_2
-from torchvision.models import mnasnet0_5
-from torchvision.models import mnasnet0_75
-from torchvision.models import mnasnet1_0
-from torchvision.models import mnasnet1_3
-from torchvision.models import resnext50_32x4d
-from torchvision.models import resnext101_32x8d
-from torchvision.models import shufflenet_v2_x1_0
-from torchvision.models import shufflenet_v2_x1_5
-from torchvision.models import shufflenet_v2_x2_0
-from torchvision.models import squeezenet1_1

predict.py

@@ -1,57 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-from torchvision import transforms
-from PIL import Image
-from models import *
-from torchmetrics import ConfusionMatrix
-import matplotlib.pyplot as plt
-from configs import *
-
-
-# Load your model (change this according to your model definition)
-MODEL.load_state_dict(
-    torch.load(MODEL_SAVE_PATH, map_location=DEVICE)
-)  # Load the model on the same device
-MODEL.eval()
-MODEL = MODEL.to(DEVICE)
-MODEL.eval()
-torch.set_grad_enabled(False)
-
-
-def predict_image(image_path, model=MODEL, transform=preprocess):
-    classes = [
-        'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease'
-    ]
-
-    print("---------------------------")
-    print("Image path:", image_path)
-    image = Image.open(image_path)
-    image = transform(image).unsqueeze(0)
-    image = image.to(DEVICE)
-    output = model(image)
-
-    # Softmax algorithm
-    probabilities = torch.softmax(output, dim=1)[0] * 100
-
-    # Sort the classes by probabilities in descending order
-    sorted_classes = sorted(
-        zip(classes, probabilities), key=lambda x: x[1], reverse=True
-    )
-
-    # Report the prediction for each class
-    print("Probabilities for each class:")
-    for class_label, class_prob in sorted_classes:
-        class_prob = class_prob.item().__round__(2)
-        print(f"{class_label}: {class_prob}%")
-
-    # Get the predicted class
-    predicted_class = sorted_classes[0][0]  # Most probable class
-    predicted_label = classes.index(predicted_class)
-
-    # Report the prediction
-    print("Predicted class:", predicted_label)
-    print("Predicted label:", predicted_class)
-    print("---------------------------")
-
-    return predicted_label, sorted_classes

train.py

@@ -1,145 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torchvision.transforms import transforms
-from torch.utils.data import DataLoader
-from torchvision.utils import make_grid
-from scipy.ndimage import gaussian_filter1d
-import matplotlib.pyplot as plt
-from models import *
-from torch.utils.tensorboard import SummaryWriter
-from configs import *
-import data_loader
-
-# Set up TensorBoard writer
-writer = SummaryWriter(log_dir="output/tensorboard/training")
-
-# Define a function for plotting and logging metrics
-def plot_and_log_metrics(metrics_dict, step, prefix="Train"):
-    for metric_name, metric_value in metrics_dict.items():
-        writer.add_scalar(f"{prefix}/{metric_name}", metric_value, step)
-
-# Data loader
-train_loader, valid_loader = data_loader.load_data(
-    RAW_DATA_DIR, AUG_DATA_DIR, EXTERNAL_DATA_DIR, preprocess
-)
-
-# Initialize model, criterion, optimizer, and scheduler
-MODEL = MODEL.to(DEVICE)
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.SGD(MODEL.parameters(), lr=LEARNING_RATE)
-scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)
-
-# Lists to store training and validation loss history
-TRAIN_LOSS_HIST = []
-VAL_LOSS_HIST = []
-AVG_TRAIN_LOSS_HIST = []
-AVG_VAL_LOSS_HIST = []
-TRAIN_ACC_HIST = []
-VAL_ACC_HIST = []
-
-# Training loop
-for epoch in range(NUM_EPOCHS):
-    MODEL.train()  # Set model to training mode
-    running_loss = 0.0
-    total_train = 0
-    correct_train = 0
-
-    for i, (inputs, labels) in enumerate(train_loader, 0):
-        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
-        optimizer.zero_grad()
-        outputs = MODEL(inputs)
-        loss = criterion(outputs, labels)
-        loss.backward()
-        optimizer.step()
-        running_loss += loss.item()
-
-        if (i + 1) % NUM_PRINT == 0:
-            print(
-                "[Epoch %d, Batch %d] Loss: %.6f"
-                % (epoch + 1, i + 1, running_loss / NUM_PRINT)
-            )
-            running_loss = 0.0
-
-        _, predicted = torch.max(outputs, 1)
-        total_train += labels.size(0)
-        correct_train += (predicted == labels).sum().item()
-
-    avg_train_loss = running_loss / len(train_loader)
-    TRAIN_LOSS_HIST.append(avg_train_loss)
-    TRAIN_ACC_HIST.append(correct_train / total_train)
-
-    # Log training metrics
-    train_metrics = {
-        "Loss": avg_train_loss,
-        "Accuracy": correct_train / total_train,
-    }
-    plot_and_log_metrics(train_metrics, epoch, prefix="Train")
-
-    # Learning rate scheduling
-    scheduler.step()
-
-    # Validation loop
-    MODEL.eval()  # Set model to evaluation mode
-    val_loss = 0.0
-    correct_val = 0
-    total_val = 0
-
-    with torch.no_grad():
-        for inputs, labels in valid_loader:
-            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
-            outputs = MODEL(inputs)
-            loss = criterion(outputs, labels)
-            val_loss += loss.item()
-            # Calculate accuracy
-            _, predicted = torch.max(outputs, 1)
-            total_val += labels.size(0)
-            correct_val += (predicted == labels).sum().item()
-
-    avg_val_loss = val_loss / len(valid_loader)
-    VAL_LOSS_HIST.append(avg_val_loss)
-    VAL_ACC_HIST.append(correct_val / total_val)
-
-    # Log validation metrics
-    val_metrics = {
-        "Loss": avg_val_loss,
-        "Accuracy": correct_val / total_val,
-    }
-    plot_and_log_metrics(val_metrics, epoch, prefix="Validation")
-
-    # Add sample images to TensorBoard
-    sample_images, _ = next(iter(valid_loader))
-    sample_images = sample_images.to(DEVICE)
-    grid_image = make_grid(
-        sample_images, nrow=8, normalize=True
-    )
-    writer.add_image("Sample Images", grid_image, global_step=epoch)
-
-# Save the model
-torch.save(MODEL.state_dict(), MODEL_SAVE_PATH)
-print("Model saved at", MODEL_SAVE_PATH)
-
-# Plot loss and accuracy curves
-plt.figure(figsize=(12, 4))
-plt.subplot(1, 2, 1)
-plt.plot(range(1, NUM_EPOCHS + 1), TRAIN_LOSS_HIST, label="Train Loss")
-plt.plot(range(1, NUM_EPOCHS + 1), VAL_LOSS_HIST, label="Validation Loss")
-plt.xlabel("Epochs")
-plt.ylabel("Loss")
-plt.legend()
-plt.title("Loss Curves")
-
-plt.subplot(1, 2, 2)
-plt.plot(range(1, NUM_EPOCHS + 1), TRAIN_ACC_HIST, label="Train Accuracy")
-plt.plot(range(1, NUM_EPOCHS + 1), VAL_ACC_HIST, label="Validation Accuracy")
-plt.xlabel("Epochs")
-plt.ylabel("Accuracy")
-plt.legend()
-plt.title("Accuracy Curves")
-
-plt.tight_layout()
-plt.savefig("training_curves.png")
-
-# Close TensorBoard writer
-writer.close()

tuning.py

@@ -1,186 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from models import *  # Import your model here
-from torch.utils.tensorboard import SummaryWriter
-from torchvision.utils import make_grid
-import optuna
-from configs import *
-import data_loader
-
-# Data loader
-train_loader, valid_loader = data_loader.load_data(
-    RAW_DATA_DIR, AUG_DATA_DIR, EXTERNAL_DATA_DIR, preprocess
-)
-
-# Initialize model, criterion, optimizer, and scheduler
-MODEL = MODEL.to(DEVICE)
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(MODEL.parameters(), lr=LEARNING_RATE)
-scheduler = optim.lr_scheduler.ReduceLROnPlateau(
-    optimizer, mode="min", factor=0.1, patience=10, verbose=True
-)
-
-# Lists to store training and validation loss history
-TRAIN_LOSS_HIST = []
-VAL_LOSS_HIST = []
-TRAIN_ACC_HIST = []
-VAL_ACC_HIST = []
-AVG_TRAIN_LOSS_HIST = []
-AVG_VAL_LOSS_HIST = []
-
-# Create a TensorBoard writer for logging
-writer = SummaryWriter(
-    log_dir="output/tensorboard/tuning",
-)
-
-# Define early stopping parameters
-early_stopping_patience = 10  # Number of epochs to wait for improvement
-best_val_loss = float('inf')
-no_improvement_count = 0
-
-def train_epoch(epoch):
-    MODEL.train(True)
-    running_loss = 0.0
-    total_train = 0
-    correct_train = 0
-
-    for i, (inputs, labels) in enumerate(train_loader, 0):
-        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
-        optimizer.zero_grad()
-        outputs = MODEL(inputs)
-        loss = criterion(outputs, labels)
-        loss.backward()
-        optimizer.step()
-        running_loss += loss.item()
-
-        if (i + 1) % NUM_PRINT == 0:
-            print(
-                "[Epoch %d, Batch %d] Loss: %.6f"
-                % (epoch + 1, i + 1, running_loss / NUM_PRINT)
-            )
-            running_loss = 0.0
-
-        _, predicted = torch.max(outputs, 1)
-        total_train += labels.size(0)
-        correct_train += (predicted == labels).sum().item()
-
-    TRAIN_LOSS_HIST.append(loss.item())
-    train_accuracy = correct_train / total_train
-    TRAIN_ACC_HIST.append(train_accuracy)
-    # Calculate the average training loss for the epoch
-    avg_train_loss = running_loss / len(train_loader)
-
-    writer.add_scalar("Loss/Train", avg_train_loss, epoch)
-    writer.add_scalar("Accuracy/Train", train_accuracy, epoch)
-    AVG_TRAIN_LOSS_HIST.append(avg_train_loss)
-
-    # Print average training loss for the epoch
-    print("[Epoch %d] Average Training Loss: %.6f" % (epoch + 1, avg_train_loss))
-
-    # Learning rate scheduling
-    lr_1 = optimizer.param_groups[0]["lr"]
-    print("Learning Rate: {:.15f}".format(lr_1))
-    scheduler.step(avg_train_loss)
-
-def validate_epoch(epoch):
-    global best_val_loss, no_improvement_count  
-    
-    MODEL.eval()
-    val_loss = 0.0
-    correct_val = 0
-    total_val = 0
-
-    with torch.no_grad():
-        for inputs, labels in valid_loader:
-            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
-            outputs = MODEL(inputs)
-            loss = criterion(outputs, labels)
-            val_loss += loss.item()
-            # Calculate accuracy
-            _, predicted = torch.max(outputs, 1)
-            total_val += labels.size(0)
-            correct_val += (predicted == labels).sum().item()
-
-    VAL_LOSS_HIST.append(loss.item())
-
-    # Calculate the average validation loss for the epoch
-    avg_val_loss = val_loss / len(valid_loader)
-    AVG_VAL_LOSS_HIST.append(loss.item())
-    print("Average Validation Loss: %.6f" % (avg_val_loss))
-
-    # Calculate the accuracy of the validation set
-    val_accuracy = correct_val / total_val
-    VAL_ACC_HIST.append(val_accuracy)
-    print("Validation Accuracy: %.6f" % (val_accuracy))
-    writer.add_scalar("Loss/Validation", avg_val_loss, epoch)
-    writer.add_scalar("Accuracy/Validation", val_accuracy, epoch)
-
-    # Add sample images to TensorBoard
-    sample_images, _ = next(iter(valid_loader))  # Get a batch of sample images
-    sample_images = sample_images.to(DEVICE)
-    grid_image = make_grid(
-        sample_images, nrow=8, normalize=True
-    )  # Create a grid of images
-    writer.add_image("Sample Images", grid_image, global_step=epoch)
-
-    # Check for early stopping
-    if avg_val_loss < best_val_loss:
-        best_val_loss = avg_val_loss
-        no_improvement_count = 0
-    else:
-        no_improvement_count += 1
-
-    if no_improvement_count >= early_stopping_patience:
-        print(f"Early stopping after {epoch + 1} epochs without improvement.")
-        return True  # Return True to stop training
-
-def objective(trial):
-    global best_val_loss, no_improvement_count
-    
-    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-1)
-    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
-
-    # Modify the model and optimizer using suggested hyperparameters
-    optimizer = optim.Adam(MODEL.parameters(), lr=learning_rate)
-
-    for epoch in range(10):
-        train_epoch(epoch)
-        early_stopping = validate_epoch(epoch)
-
-        # Check for early stopping
-        if early_stopping:
-            break
-
-    # Calculate a weighted score based on validation accuracy and loss
-    validation_score = VAL_ACC_HIST[-1] - AVG_VAL_LOSS_HIST[-1]
-
-    # Return the negative score as Optuna maximizes by default
-    return -validation_score
-
-if __name__ == "__main__":
-    study = optuna.create_study(direction="maximize")
-    study.optimize(objective, n_trials=100, show_progress_bar=True)
-
-    # Print statistics
-    print("Number of finished trials: ", len(study.trials))
-    pruned_trials = [
-        t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED
-    ]
-    print("Number of pruned trials: ", len(pruned_trials))
-    complete_trials = [
-        t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE
-    ]
-    print("Number of complete trials: ", len(complete_trials))
-
-    # Print best trial
-    trial = study.best_trial
-    print("Best trial:")
-    print("  Value: ", -trial.value)  # Negate the value as it was maximized
-    print("  Params: ")
-    for key, value in trial.params.items():
-        print(f"    {key}: {value}")
-
-    # Close TensorBoard writer
-    writer.close()