Update

augment.py

@@ -2,30 +2,33 @@ import os
 import Augmentor
 import shutil
 from configs import *
+import uuid
 
 tasks = ["1", "2", "3", "4", "5", "6"]
 
-for task in tasks:
+for task in ["1"]:
     # 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)
+    for class_label in ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']:
+        if class_label != ".DS_Store":
+            print("Augmenting images in class: ", class_label, " 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}",
-                )
+            if not os.path.exists(f"{TEMP_DATA_DIR}Task {task}/{class_label}/"):
+                os.makedirs(f"{TEMP_DATA_DIR}Task {task}/{class_label}/")
+            if os.path.exists(f"{RAW_DATA_DIR}Task {task}/{class_label}"):
+                for file in os.listdir(f"{RAW_DATA_DIR}Task {task}/{class_label}"):
+                    shutil.copy(
+                        f"{RAW_DATA_DIR}Task {task}/{class_label}/{file}",
+                        f"{TEMP_DATA_DIR}Task {task}/{class_label}/{str(uuid.uuid4())}.png",
+                    )
+            if os.path.exists(f"{EXTERNAL_DATA_DIR}Task {task}/{class_label}"):
+                for file in os.listdir(f"{EXTERNAL_DATA_DIR}Task {task}/{class_label}"):
+                    shutil.copy(
+                        f"{EXTERNAL_DATA_DIR}Task {task}/{class_label}/{file}",
+                        f"{TEMP_DATA_DIR}Task {task}/{class_label}/{str(uuid.uuid4())}.png",
+                    )
             p = Augmentor.Pipeline(
-                f"data/temp/Task {task}/{disease}",
-                output_directory=f"{disease}/",
+                f"{TEMP_DATA_DIR}Task {task}/{class_label}",
+                output_directory=f"{class_label}/",
                 save_format="png",
             )
             p.rotate(probability=0.8, max_left_rotation=5, max_right_rotation=5)
@@ -39,20 +42,20 @@ for task in tasks:
             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}/")
+            if not os.path.exists(f"{AUG_DATA_DIR}Task {task}/"):
+                os.makedirs(f"{AUG_DATA_DIR}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}",
+                f"{TEMP_DATA_DIR}Task {task}/{class_label}/{class_label}",
+                f"{AUG_DATA_DIR}Task {task}/{class_label}",
             )
             # Rename all the augmented images to [01, 02, 03]
             number = 0
-            for file in os.listdir(f"data/train/augmented/Task {task}/{disease}"):
+            for file in os.listdir(f"{AUG_DATA_DIR}Task {task}/{class_label}"):
                 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",
+                    f"{AUG_DATA_DIR}Task {task}/{class_label}/{file}",
+                    f"{AUG_DATA_DIR}Task {task}/{class_label}/{number}.png",
                 )

configs.py

@@ -6,23 +6,23 @@ from models import *
 
 # Constants
 RANDOM_SEED = 123
-BATCH_SIZE = 64
+BATCH_SIZE = 16
 NUM_EPOCHS = 100
-LEARNING_RATE = 1.6317268278715415e-05
-OPTIMIZER_NAME = "Adam"
+LEARNING_RATE = 5.847227637580824e-05
 STEP_SIZE = 10
-GAMMA = 0.5
+GAMMA = 1.0
 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)
+RAW_DATA_DIR = r"data/train/raw/Task "
+AUG_DATA_DIR = r"data/train/augmented/Task "
+EXTERNAL_DATA_DIR = r"data/train/external/Task "
+TEMP_DATA_DIR = "data/temp/"
 NUM_CLASSES = 7
-# Define classes as listdir of augmented data
-CLASSES = os.listdir("data/train/augmented/Task 1/")
+EARLY_STOPPING_PATIENCE = 20
+CLASSES = ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']
 MODEL_SAVE_PATH = "output/checkpoints/model.pth"
-MODEL = googlenet(num_classes=NUM_CLASSES)
+MODEL = efficientnet_b1(num_classes=NUM_CLASSES)
 
 print(CLASSES)
 

eval.py

@@ -1,10 +1,9 @@
 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
+from torchmetrics import ConfusionMatrix, Accuracy, F1Score
 import matplotlib.pyplot as plt
 from configs import *
 from data_loader import load_data  # Import the load_data function
@@ -19,7 +18,6 @@ MODEL = MODEL.to(DEVICE)
 MODEL.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=DEVICE))
 MODEL.eval()
 
-
 def predict_image(image_path, model, transform):
     model.eval()
     correct_predictions = 0
@@ -32,6 +30,9 @@ def predict_image(image_path, model, transform):
     true_classes = []
     predicted_labels = []
 
+    accuracy_metric = Accuracy(num_classes=NUM_CLASSES, task="multiclass")
+    f1_metric = F1Score(num_classes=NUM_CLASSES, task="multiclass")
+
     with torch.no_grad():
         for image_file in images:
             print("---------------------------")
@@ -57,7 +58,7 @@ def predict_image(image_path, model, transform):
     # Calculate accuracy and f1 score
     accuracy = correct_predictions / total_predictions
     print("Accuracy:", accuracy)
-    f1 = f1_score(true_classes, predicted_labels, average="weighted")
+    f1 = f1_metric(torch.tensor(predicted_labels), torch.tensor(true_classes)).item()
     print("Weighted F1 Score:", f1)
 
     # Convert the lists to tensors
@@ -66,13 +67,12 @@ def predict_image(image_path, model, transform):
 
     # Create a confusion matrix
     conf_matrix = ConfusionMatrix(num_classes=NUM_CLASSES, task="multiclass")
-    conf_matrix.update(predicted_labels_tensor, true_classes_tensor)
+    conf_matrix(predicted_labels_tensor, true_classes_tensor)
 
     # Plot the confusion matrix
     conf_matrix.compute()
     conf_matrix.plot()
     plt.show()
 
-
 # Call predict_image function
 predict_image(image_path, MODEL, preprocess)

models.py

@@ -34,3 +34,9 @@ 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
+from torchvision.models import efficientnet_v2_s
+from torchvision.models import efficientnet_v2_m
+from torchvision.models import efficientnet_v2_l
+from torchvision.models import efficientnet_b0
+from torchvision.models import efficientnet_b1
+

train.py

@@ -8,50 +8,32 @@ 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"):
+def setup_tensorboard():
+    return SummaryWriter(log_dir="output/tensorboard/training")
+
+
+def load_and_preprocess_data():
+    return data_loader.load_data(
+        RAW_DATA_DIR + str(TASK), AUG_DATA_DIR + str(TASK), EXTERNAL_DATA_DIR + str(TASK), preprocess
+    )
+
+
+def initialize_model_optimizer_scheduler():
+    model = MODEL.to(DEVICE)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)
+    return model, criterion, optimizer, scheduler
+
+
+def plot_and_log_metrics(metrics_dict, step, writer, 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()
-if OPTIMIZER_NAME == "LBFGS":
-    optimizer = optim.LBFGS(MODEL.parameters(), lr=LEARNING_RATE)
-elif OPTIMIZER_NAME == "Adam":
-    optimizer = optim.Adam(MODEL.parameters(), lr=LEARNING_RATE)
-elif OPTIMIZER_NAME == "SGD":
-    optimizer = optim.SGD(MODEL.parameters(), lr=LEARNING_RATE)
-    
-scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)
-
-# Define early stopping parameters
-early_stopping_patience = 20  # Number of epochs with no improvement to wait before stopping
-best_val_loss = float("inf")
-best_val_accuracy = 0.0
-no_improvement_count = 0
-
-# 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):
-    print(f"[Epoch: {epoch + 1}]")
-    print("Learning rate:", scheduler.get_last_lr()[0])
-    MODEL.train()  # Set model to training mode
+
+def train_one_epoch(model, criterion, optimizer, train_loader, epoch):
+    model.train()
     running_loss = 0.0
     total_train = 0
     correct_train = 0
@@ -59,16 +41,13 @@ for epoch in range(NUM_EPOCHS):
     for i, (inputs, labels) in enumerate(train_loader, 0):
         inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
         optimizer.zero_grad()
-        if MODEL.__class__.__name__ == "GoogLeNet": # the shit GoogLeNet has a different output
-            outputs = MODEL(inputs).logits
+        if model.__class__.__name__ == "GoogLeNet":
+            outputs = model(inputs).logits
         else:
-            outputs = MODEL(inputs)
+            outputs = model(inputs)
         loss = criterion(outputs, labels)
         loss.backward()
-        if OPTIMIZER_NAME == "LBFGS":
-            optimizer.step(closure=lambda: loss)
-        else:
-            optimizer.step()
+        optimizer.step()
         running_loss += loss.item()
 
         if (i + 1) % NUM_PRINT == 0:
@@ -83,21 +62,11 @@ for epoch in range(NUM_EPOCHS):
         correct_train += (predicted == labels).sum().item()
 
     avg_train_loss = running_loss / len(train_loader)
-    AVG_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
+    return avg_train_loss, correct_train / total_train
+
+
+def validate_model(model, criterion, valid_loader):
+    model.eval()
     val_loss = 0.0
     correct_val = 0
     total_val = 0
@@ -105,67 +74,120 @@ for epoch in range(NUM_EPOCHS):
     with torch.no_grad():
         for inputs, labels in valid_loader:
             inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
-            outputs = MODEL(inputs)
+            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)
-    AVG_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")
-    
-    # Print average training and validation metrics
-    print(f"Average Training Loss: {avg_train_loss:.6f}")
-    print(f"Average Validation Loss: {avg_val_loss:.6f}")
-    print(f"Training Accuracy: {correct_train / total_train:.6f}")
-    print(f"Validation Accuracy: {correct_val / total_val:.6f}")
-    
-    # Check for early stopping based on validation accuracy
-    if correct_val / total_val > best_val_accuracy:
-        best_val_accuracy = correct_val / total_val
-        no_improvement_count = 0
-    else:
-        no_improvement_count += 1
-
-    # Early stopping condition
-    if no_improvement_count >= early_stopping_patience:
-        print("Early stopping: Validation accuracy did not improve for {} consecutive epochs.".format(early_stopping_patience))
-        break  # Stop training
-
-# 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, len(AVG_TRAIN_LOSS_HIST) + 1), AVG_TRAIN_LOSS_HIST, label="Average Train Loss")
-plt.plot(range(1, len(AVG_VAL_LOSS_HIST) + 1), AVG_VAL_LOSS_HIST, label="Average Validation Loss")
-plt.xlabel("Epochs")
-plt.ylabel("Loss")
-plt.legend()
-plt.title("Loss Curves")
-
-plt.subplot(1, 2, 2)
-plt.plot(range(1, len(TRAIN_ACC_HIST) + 1), TRAIN_ACC_HIST, label="Train Accuracy")
-plt.plot(range(1, len(VAL_ACC_HIST) + 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()
+    return avg_val_loss, correct_val / total_val
+
+
+def main_training_loop():
+    writer = setup_tensorboard()
+    train_loader, valid_loader = load_and_preprocess_data()
+    model, criterion, optimizer, scheduler = initialize_model_optimizer_scheduler()
+
+    best_val_loss = float("inf")
+    best_val_accuracy = 0.0
+    no_improvement_count = 0
+
+    AVG_TRAIN_LOSS_HIST = []
+    AVG_VAL_LOSS_HIST = []
+    TRAIN_ACC_HIST = []
+    VAL_ACC_HIST = []
+
+    for epoch in range(NUM_EPOCHS):
+        print(f"[Epoch: {epoch + 1}]")
+        print("Learning rate:", scheduler.get_last_lr()[0])
+
+        avg_train_loss, train_accuracy = train_one_epoch(
+            model, criterion, optimizer, train_loader, epoch
+        )
+        AVG_TRAIN_LOSS_HIST.append(avg_train_loss)
+        TRAIN_ACC_HIST.append(train_accuracy)
+
+        # Log training metrics
+        train_metrics = {
+            "Loss": avg_train_loss,
+            "Accuracy": train_accuracy,
+        }
+        plot_and_log_metrics(train_metrics, epoch, writer=writer, prefix="Train")
+
+        # Learning rate scheduling
+        scheduler.step()
+
+        avg_val_loss, val_accuracy = validate_model(model, criterion, valid_loader)
+        AVG_VAL_LOSS_HIST.append(avg_val_loss)
+        VAL_ACC_HIST.append(val_accuracy)
+
+        # Log validation metrics
+        val_metrics = {
+            "Loss": avg_val_loss,
+            "Accuracy": val_accuracy,
+        }
+        plot_and_log_metrics(train_metrics, epoch, writer=writer, prefix="Train")
+
+        # Print average training and validation metrics
+        print(f"Average Training Loss: {avg_train_loss:.6f}")
+        print(f"Average Validation Loss: {avg_val_loss:.6f}")
+        print(f"Training Accuracy: {train_accuracy:.6f}")
+        print(f"Validation Accuracy: {val_accuracy:.6f}")
+
+        # Check for early stopping based on validation accuracy
+        if val_accuracy > best_val_accuracy:
+            best_val_accuracy = val_accuracy
+            no_improvement_count = 0
+        else:
+            no_improvement_count += 1
+
+        # Early stopping condition
+        if no_improvement_count >= EARLY_STOPPING_PATIENCE:
+            print(
+                "Early stopping: Validation accuracy did not improve for {} consecutive epochs.".format(
+                    EARLY_STOPPING_PATIENCE
+                )
+            )
+            break
+
+    # 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, len(AVG_TRAIN_LOSS_HIST) + 1),
+        AVG_TRAIN_LOSS_HIST,
+        label="Average Train Loss",
+    )
+    plt.plot(
+        range(1, len(AVG_VAL_LOSS_HIST) + 1),
+        AVG_VAL_LOSS_HIST,
+        label="Average Validation Loss",
+    )
+    plt.xlabel("Epochs")
+    plt.ylabel("Loss")
+    plt.legend()
+    plt.title("Loss Curves")
+
+    plt.subplot(1, 2, 2)
+    plt.plot(range(1, len(TRAIN_ACC_HIST) + 1), TRAIN_ACC_HIST, label="Train Accuracy")
+    plt.plot(range(1, len(VAL_ACC_HIST) + 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()
+
+
+if __name__ == "__main__":
+    main_training_loop()

tuning.py

@@ -9,21 +9,27 @@ from configs import *
 import data_loader
 from torch.utils.tensorboard import SummaryWriter
 
-optuna.logging.set_verbosity(optuna.logging.DEBUG)
-
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 EPOCHS = 10
+N_TRIALS = 50
+TIMEOUT = 3600  # 1 hour
 
 # Create a TensorBoard writer
-writer = SummaryWriter(log_dir="output/tensorboard/tuning/", )
+writer = SummaryWriter(log_dir="output/tensorboard/tuning")
+
 
 def create_data_loaders(batch_size):
     # Create or modify data loaders with the specified batch size
     train_loader, valid_loader = data_loader.load_data(
-        RAW_DATA_DIR, AUG_DATA_DIR, EXTERNAL_DATA_DIR, preprocess, batch_size=batch_size
+        RAW_DATA_DIR + str(TASK),
+        AUG_DATA_DIR + str(TASK),
+        EXTERNAL_DATA_DIR + str(TASK),
+        preprocess,
+        batch_size=batch_size,
     )
     return train_loader, valid_loader
 
+
 def objective(trial, model=MODEL):
     # Generate the model.
     model = model.to(DEVICE)
@@ -35,11 +41,16 @@ def objective(trial, model=MODEL):
     train_loader, valid_loader = create_data_loaders(batch_size)
 
     # Generate the optimizer.
-    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "SGD"])
-    lr = trial.suggest_float("lr", 1e-5, 1e-3, log=True)
-    optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
+    lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
+    optimizer = optim.Adam(model.parameters(), lr=lr)
     criterion = nn.CrossEntropyLoss()
 
+    # Suggest the gamma parameter for the learning rate scheduler.
+    gamma = trial.suggest_float("gamma", 0.1, 1.0, step=0.1)
+
+    # Create a learning rate scheduler with the suggested gamma.
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma)
+
     # Training of the model.
     for epoch in range(EPOCHS):
         print(f"[Epoch: {epoch} | Trial: {trial.number}]")
@@ -47,16 +58,18 @@ def objective(trial, model=MODEL):
         for batch_idx, (data, target) in enumerate(train_loader, 0):
             data, target = data.to(DEVICE), target.to(DEVICE)
             optimizer.zero_grad()
-            if model.__class__.__name__ == "GoogLeNet": # the shit GoogLeNet has a different output
+            if (
+                model.__class__.__name__ == "GoogLeNet"
+            ):  # the shit GoogLeNet has a different output
                 output = model(data).logits
             else:
                 output = model(data)
             loss = criterion(output, target)
             loss.backward()
-            if optimizer_name == "LBFGS":
-                optimizer.step(closure=lambda: loss)
-            else:
-                optimizer.step()
+            optimizer.step()
+
+        # Update the learning rate using the scheduler.
+        scheduler.step()
 
         # Validation of the model.
         model.eval()
@@ -74,14 +87,8 @@ def objective(trial, model=MODEL):
         # Log hyperparameters and accuracy to TensorBoard
         writer.add_scalar("Accuracy", accuracy, trial.number)
         writer.add_hparams(
-            {
-                "batch_size": batch_size,
-                "optimizer": optimizer_name,
-                "lr": lr
-            },
-            {
-                "accuracy": accuracy
-            }
+            {"batch_size": batch_size, "lr": lr, "gamma": gamma},
+            {"accuracy": accuracy},
         )
 
         # Print hyperparameters and accuracy
@@ -93,29 +100,29 @@ def objective(trial, model=MODEL):
         if trial.should_prune():
             raise optuna.exceptions.TrialPruned()
 
+    if trial.number > 10 and trial.params["lr"] < 1e-3 and accuracy < 0.7:
+        return float("inf")  # Prune the trial
+
     return accuracy
 
+
 if __name__ == "__main__":
     pruner = optuna.pruners.HyperbandPruner()
-    study = optuna.create_study(direction="maximize", pruner=pruner, study_name="handetect")
-    study.optimize(objective, n_trials=100, timeout=1000)
-
-    pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED])
-    complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE])
+    study = optuna.create_study(
+        direction="maximize",  # Adjust the direction as per your optimization goal
+        pruner=pruner,
+        study_name="hyperparameter_tuning",
+    )
 
-    print("Study statistics: ")
-    print("  Number of finished trials: ", len(study.trials))
-    print("  Number of pruned trials: ", len(pruned_trials))
-    print("  Number of complete trials: ", len(complete_trials))
+    # Optimize the hyperparameters
+    study.optimize(
+        objective, n_trials=100, timeout=3600
+    )  # Adjust the number of trials and timeout as needed
 
+    # Print the best trial
+    best_trial = study.best_trial
     print("Best trial:")
-    trial = study.best_trial
-
-    print("  Value: ", trial.value)
-
+    print("  Value: ", best_trial.value)
     print("  Params: ")
-    for key, value in trial.params.items():
+    for key, value in best_trial.params.items():
         print("    {}: {}".format(key, value))
-
-    # Close TensorBoard writer
-    writer.close()