Update

handetect/augment.py

@@ -0,0 +1,41 @@
+import os
+import Augmentor
+
+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 i in os.listdir("data/train/raw/Task " + task):
+        if i != ".DS_Store":
+            print("Augmenting images in class: ", i)
+            p = Augmentor.Pipeline(f"data/train/raw/Task {task}/{i}", output_directory=i, 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/train/raw/Task {task}/{i}/{i}",
+                f"data/train/augmented/Task {task}/{i}",
+            )
+            # Rename all the augmented images to [01, 02, 03]
+            number = 0
+            for j in os.listdir(f"data/train/augmented/Task {task}/{i}"):
+                number = int(number) + 1
+                if len(str(number)) == 1:
+                    number = "0" + str(number)
+                os.rename(
+                    f"data/train/augmented/Task {task}/{i}/{j}",
+                    f"data/train/augmented/Task {task}/{i}/{number}.png",
+                )
+
+

handetect/configs.py

@@ -0,0 +1,41 @@
+import os
+import torch
+from torchvision import transforms
+from torch.utils.data import Dataset
+from models import *
+
+# Constants
+RANDOM_SEED = 123
+BATCH_SIZE = 64
+NUM_EPOCHS = 100
+LEARNING_RATE = 0.02750299610194638
+STEP_SIZE = 10
+GAMMA = 0.5
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+NUM_PRINT = 100
+TASK = 1
+ORIG_DATA_DIR = r"data/train/raw/Task " + str(TASK)
+AUG_DATA_DIR = r"data/train/augmented/Task " + str(TASK)
+NUM_CLASSES = len(os.listdir(ORIG_DATA_DIR))
+MODEL_SAVE_PATH = "output/checkpoints/model.pth"
+MODEL = shufflenet_v2_x0_5(num_classes=NUM_CLASSES)
+
+preprocess = transforms.Compose(
+    [
+        transforms.Resize((64, 64)),  # Resize images to 64x64
+        transforms.ToTensor(),  # Convert to tensor
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize
+    ]
+)
+
+# 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

handetect/data_loader.py

@@ -0,0 +1,30 @@
+from configs import *
+from torchvision.datasets import ImageFolder
+from torch.utils.data import random_split, DataLoader, Dataset
+
+
+def load_data(original_dir, augmented_dir, preprocess):
+    # Load the dataset using ImageFolder
+    original_dataset = ImageFolder(root=original_dir, transform=preprocess)
+    augmented_dataset = ImageFolder(root=augmented_dir, transform=preprocess)
+    dataset = original_dataset + augmented_dataset
+
+    print("Classes: ", *original_dataset.classes, sep = ' ')
+    print("Length of original dataset: ", len(original_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

handetect/eval.py

@@ -0,0 +1,91 @@
+import os
+import torch
+from torchvision.transforms import transforms
+from sklearn.metrics import f1_score
+from models import * 
+import pathlib
+from PIL import Image
+from torchmetrics import ConfusionMatrix, Accuracy
+import matplotlib.pyplot as plt
+from configs import *
+
+image_path = "data/test/Task 1/" 
+
+# constants
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+NUM_CLASSES = 6 
+
+# load the model
+images = list(pathlib.Path(image_path).rglob("*.png"))
+classes = os.listdir(image_path)
+print(images)
+
+true_classs = []
+predicted_labels = []
+
+model = mobilenet_v3_small(pretrained=False, num_classes=NUM_CLASSES)  
+model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=DEVICE)) 
+model.eval()
+model = model.to(DEVICE)
+
+
+# Define transformation for preprocessing
+preprocess = transforms.Compose(
+    [
+        transforms.Resize((64, 64)),  # Resize images to 64x64
+        transforms.Grayscale(num_output_channels=3),  # Convert to grayscale
+        transforms.ToTensor(),  # Convert to tensor
+        transforms.Normalize((0.5,), (0.5,)),  # Normalize (for grayscale)
+    ]
+)
+
+# evaluate the model
+all_predictions = []
+true_labels = []
+
+def predict_image(image_path, model, transform):
+    model.eval()
+    correct_predictions = 0
+    total_predictions = len(images)
+    
+    with torch.no_grad():
+        for i in images:
+            print('---------------------------')
+            # Check the true label of the image by checking the sequence of the folder in Task 1
+            true_class = classes.index(i.parts[-2])
+            print("Image path:", i)
+            print("True class:", true_class)
+            image = Image.open(i)
+            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_classs.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 socre
+    accuracy = correct_predictions / total_predictions
+    print("Accuracy:", accuracy)
+    f1 = f1_score(true_classs, predicted_labels, average='weighted')
+    print("Weighted F1 Score:", f1)
+
+# Call predict_image function
+predict_image(image_path, model, preprocess)
+
+# Convert the lists to tensors
+predicted_labels_tensor = torch.tensor(predicted_labels)
+true_classs_tensor = torch.tensor(true_classs)
+
+conf_matrix = ConfusionMatrix(num_classes=NUM_CLASSES, task='multiclass')
+conf_matrix.update(predicted_labels_tensor, true_classs_tensor)
+
+# Plot confusion matrix
+conf_matrix.plot()
+plt.show()

handetect/models.py

@@ -2,7 +2,7 @@
 # This file stores all the models used in the project.#
 #######################################################
 
-import torch
+# Import all models from torchvision.models
 from torchvision.models import resnet50
 from torchvision.models import resnet18
 from torchvision.models import squeezenet1_0
@@ -13,3 +13,5 @@ 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

handetect/predict.py

@@ -6,36 +6,20 @@ from PIL import Image
 from handetect.models import *
 from torchmetrics import ConfusionMatrix
 import matplotlib.pyplot as plt
+from configs import *
 
-# Define the path to your model checkpoint
-model_checkpoint_path = "model.pth"
-
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-NUM_CLASSES = 6
-
-# Define transformation for preprocessing the input image
-preprocess = transforms.Compose(
-    [
-        transforms.Resize((64, 64)),  # Resize the image to match training input size
-        transforms.Grayscale(num_output_channels=3),  # Convert the image to grayscale
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize the image
-    ]
-)
 
 # Load your model (change this according to your model definition)
-model = mobilenet_v3_small(pretrained=False, num_classes=NUM_CLASSES)
-model.load_state_dict(
-    torch.load(model_checkpoint_path, map_location=DEVICE)
+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()
+MODEL.eval()
+MODEL = MODEL.to(DEVICE)
+MODEL.eval()
 torch.set_grad_enabled(False)
 
 
-def predict_image(image_path, model=model, transform=preprocess):
+def predict_image(image_path, model=MODEL, transform=preprocess):
     # Define images variable to recursively list all the data file in the image_path
     classes = ['Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']
 

handetect/train.py

@@ -0,0 +1,174 @@
+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, random_split, Dataset
+from torchvision.datasets import ImageFolder
+import matplotlib.pyplot as plt
+from models import *
+from scipy.ndimage import gaussian_filter1d
+from torch.utils.tensorboard import SummaryWriter  # print to tensorboard
+from torchvision.utils import make_grid
+import pandas as pd
+from configs import *
+import data_loader
+
+# torch.cuda.empty_cache()
+# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:1024"
+
+writer = SummaryWriter(log_dir="output/tensorboard")
+
+
+# Data loader
+train_loader, valid_loader = data_loader.load_data(
+    ORIG_DATA_DIR, AUG_DATA_DIR, preprocess
+)
+
+
+# Initialize model, criterion, optimizer, and scheduler
+MODEL = MODEL.to(DEVICE)
+criterion = nn.CrossEntropyLoss()
+# Adam optimizer
+optimizer = optim.Adam(MODEL.parameters(), lr=LEARNING_RATE)
+# StepLR scheduler
+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(True)  # 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()
+
+    TRAIN_ACC_HIST.append(correct_train / total_train)
+
+    TRAIN_LOSS_HIST.append(loss.item())
+
+    # 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", correct_train / total_train, 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()
+
+    # 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()
+
+    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 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)
+
+# End of training loop
+
+# Save the model
+
+torch.save(MODEL.state_dict(), MODEL_SAVE_PATH)
+print("Model saved at", MODEL_SAVE_PATH)
+
+print("Generating loss plot...")
+# train_loss_line = gaussian_filter1d(TRAIN_LOSS_HIST, sigma=10)
+# val_loss_line = gaussian_filter1d(VAL_LOSS_HIST, sigma=10)
+# plt.plot(range(1, NUM_EPOCHS + 1), train_loss_line, label='Train Loss')
+# plt.plot(range(1, NUM_EPOCHS + 1), val_loss_line, label='Validation Loss')
+avg_train_loss_line = gaussian_filter1d(AVG_TRAIN_LOSS_HIST, sigma=2)
+avg_val_loss_line = gaussian_filter1d(AVG_VAL_LOSS_HIST, sigma=2)
+train_loss_line = gaussian_filter1d(TRAIN_LOSS_HIST, sigma=2)
+val_loss_line = gaussian_filter1d(VAL_LOSS_HIST, sigma=2)
+train_acc_line = gaussian_filter1d(TRAIN_ACC_HIST, sigma=2)
+val_acc_line = gaussian_filter1d(VAL_ACC_HIST, sigma=2)
+plt.plot(range(1, NUM_EPOCHS + 1), train_loss_line, label="Train Loss")
+plt.plot(range(1, NUM_EPOCHS + 1), val_loss_line, label="Validation Loss")
+plt.xlabel("Epochs")
+plt.ylabel("Loss")
+plt.legend()
+plt.title("Train Loss and Validation Loss")
+plt.savefig("loss_plot.png")
+plt.clf()
+plt.plot(range(1, NUM_EPOCHS + 1), avg_train_loss_line, label="Average Train Loss")
+plt.plot(range(1, NUM_EPOCHS + 1), avg_val_loss_line, label="Average Validation Loss")
+plt.xlabel("Epochs")
+plt.ylabel("Loss")
+plt.legend()
+plt.title("Average Train Loss and Average Validation Loss")
+plt.savefig("avg_loss_plot.png")
+plt.clf()
+plt.plot(range(1, NUM_EPOCHS + 1), train_acc_line, label="Train Accuracy")
+plt.plot(range(1, NUM_EPOCHS + 1), val_acc_line, label="Validation Accuracy")
+plt.xlabel("Epochs")
+plt.ylabel("Accuracy")
+plt.legend()
+plt.title("Train Accuracy and Validation Accuracy")
+plt.savefig("accuracy_plot.png")
+
+dummy_input = torch.randn(1, 3, 64, 64).to(DEVICE)  # Adjust input shape accordingly
+writer.add_graph(MODEL, dummy_input)
+# Close TensorBoard writer
+writer.close()

handetect/tuning.py

@@ -0,0 +1,203 @@
+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, random_split, Dataset
+from torchvision.datasets import ImageFolder
+from models import *
+from torch.utils.tensorboard import SummaryWriter #print to tensorboard
+from torchvision.utils import make_grid
+import optuna
+from configs import *
+
+writer = SummaryWriter()
+
+# Error if the classes in the original dataset and augmented dataset are not the same
+assert (
+    os.listdir(ORIG_DATA_DIR) == os.listdir(AUG_DATA_DIR)
+), "Classes in original dataset and augmented dataset are not the same"
+
+
+# Load the dataset using ImageFolder
+original_dataset = ImageFolder(root=ORIG_DATA_DIR, transform=preprocess)
+augmented_dataset = ImageFolder(root=AUG_DATA_DIR, transform=preprocess)
+dataset = original_dataset + augmented_dataset
+
+print("Classes: ", original_dataset.classes)
+print("Length of original dataset: ", len(original_dataset))
+print("Length of augmented dataset: ", len(augmented_dataset))
+print("Length of total dataset: ", len(dataset))
+
+
+# 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
+
+
+# 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
+)
+
+# Initialize model, criterion, optimizer, and scheduler
+MODEL = MODEL.to(DEVICE)
+criterion = nn.CrossEntropyLoss()
+# Adam optimizer
+optimizer = optim.Adam(MODEL.parameters(), lr=LEARNING_RATE)
+
+# ReduceLROnPlateau scheduler
+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 = []
+AVG_TRAIN_LOSS_HIST = []
+AVG_VAL_LOSS_HIST = []
+TRAIN_ACC_HIST = []
+VAL_ACC_HIST = []
+
+def objective(trial):
+    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-1, log=True)
+    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(NUM_EPOCHS):
+        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)
+
+    # 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()
+
+    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 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)
+    # Validation loop
+    MODEL.eval()  # Set model to evaluation mode
+    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)
+            _, predicted = torch.max(outputs, 1)
+            total_val += labels.size(0)
+            correct_val += (predicted == labels).sum().item()
+
+    # suan evaluation score 
+    evaluation_score = correct_val / total_val
+
+    # Return the evaluation score 
+    return evaluation_score
+
+
+if __name__ == "__main__":
+    study = optuna.create_study(direction="maximize")
+    study.optimize(objective, n_trials=300, timeout=800)
+
+    # 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)
+    print("  Params: ")
+    for key, value in trial.params.items():
+        print(f"    {key}: {value}")
+
+