Update

app.py

@@ -30,7 +30,7 @@ def process_file(webcam_filepath, upload_filepath):
 
 
 demo = gr.Interface(
-    theme='gradio/soft',
+    theme="gradio/soft",
     fn=process_file,
     title="HANDETECT",
     description="An innovative AI-powered system that facilitates early detection and monitoring of movement disorders through handwriting assessment",

augment.py

@@ -8,7 +8,15 @@ tasks = ["1", "2", "3", "4", "5", "6"]
 
 for task in ["1"]:
     # Loop through all folders in Task 1 and generate augmented images for each class
-    for class_label in ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']:
+    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

configs.py

@@ -4,15 +4,21 @@ from torchvision import transforms
 from torch.utils.data import Dataset
 from models import *
 import torch.nn as nn
-from torchvision.models import squeezenet1_0, SqueezeNet1_0_Weights
+from torchvision.models import (
+    squeezenet1_0,
+    SqueezeNet1_0_Weights,
+    mobilenet_v3_small,
+    MobileNet_V3_Small_Weights,
+)
 from torchvision.models import squeezenet1_0
+
 # Constants
 RANDOM_SEED = 123
-BATCH_SIZE = 16
+BATCH_SIZE = 32
 NUM_EPOCHS = 40
-LEARNING_RATE = 5.488903014780378e-05
+LEARNING_RATE = 0.005873024218838728
 STEP_SIZE = 10
-GAMMA = 0.3
+GAMMA = 0.9
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 NUM_PRINT = 100
 TASK = 1
@@ -22,21 +28,28 @@ EXTERNAL_DATA_DIR = r"data/train/external/Task "
 TEMP_DATA_DIR = "data/temp/"
 NUM_CLASSES = 7
 EARLY_STOPPING_PATIENCE = 20
-CLASSES = ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']
+CLASSES = [
+    "Alzheimer Disease",
+    "Cerebral Palsy",
+    "Dystonia",
+    "Essential Tremor",
+    "Healthy",
+    "Huntington Disease",
+    "Parkinson Disease",
+]
 MODEL_SAVE_PATH = "output/checkpoints/model.pth"
 
 
-    
 class SqueezeNet1_0WithDropout(nn.Module):
     def __init__(self, num_classes=1000):
         super(SqueezeNet1_0WithDropout, self).__init__()
-        squeezenet = squeezenet1_0(weights=SqueezeNet1_0_Weights) 
+        squeezenet = squeezenet1_0(weights=SqueezeNet1_0_Weights.DEFAULT)
         self.features = squeezenet.features
         self.classifier = nn.Sequential(
             nn.Conv2d(512, num_classes, kernel_size=1),
             nn.BatchNorm2d(num_classes),  # add batch normalization
             nn.ReLU(inplace=True),
-            nn.AdaptiveAvgPool2d((1, 1))
+            nn.AdaptiveAvgPool2d((1, 1)),
         )
 
     def forward(self, x):
@@ -44,9 +57,67 @@ class SqueezeNet1_0WithDropout(nn.Module):
         x = self.classifier(x)
         x = torch.flatten(x, 1)
         return x
-    
-    
-    
+
+
+# class ShuffleNetV2WithDropout(nn.Module):
+#     def __init__(self, num_classes=1000):
+#         super(ShuffleNetV2WithDropout, self).__init__()
+#         shufflenet = shufflenet_v2_x2_0(weights=ShuffleNet_V2_X2_0_Weights)
+#         self.features = shufflenet.features
+#         self.classifier = nn.Sequential(
+#             nn.Conv2d(1024, num_classes, kernel_size=1),
+#             nn.BatchNorm2d(num_classes),  # add batch normalization
+#             nn.ReLU(inplace=True),
+#             nn.AdaptiveAvgPool2d((1, 1))
+#         )
+
+#     def forward(self, x):
+#         x = self.features(x)
+#         x = self.classifier(x)
+#         x = torch.flatten(x, 1)
+#         return x
+
+
+class MobileNetV3SmallWithDropout(nn.Module):
+    def __init__(self, num_classes=1000):
+        super(MobileNetV3SmallWithDropout, self).__init__()
+        mobilenet = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights)
+        self.features = mobilenet.features
+        self.classifier = nn.Sequential(
+            nn.Conv2d(576, num_classes, kernel_size=1),
+            nn.BatchNorm2d(num_classes),  # add batch normalization
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        x = torch.flatten(x, 1)
+        return x
+
+
+class ResNet18WithNorm(nn.Module):
+    def __init__(self, num_classes=1000):
+        super(ResNet18WithNorm, self).__init__()
+        resnet = resnet18(pretrained=False)
+        self.features = nn.Sequential(
+            *list(resnet.children())[:-2]
+        )  # Remove last 2 layers (avgpool and fc)
+        self.classifier = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(),
+            nn.Linear(512, num_classes),
+            nn.BatchNorm2d(num_classes),  # Add batch normalization
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        x = torch.flatten(x, 1)
+        return x
+
+
 MODEL = SqueezeNet1_0WithDropout(num_classes=7)
 print(CLASSES)
 
@@ -70,4 +141,4 @@ class CustomDataset(Dataset):
 
     def __getitem__(self, idx):
         img, label = self.data[idx]
-        return img, label
+        return img, label

convert.py

@@ -2,6 +2,16 @@ import torch
 import onnx2tf
 from configs import *
 
-torch.onnx.export(model=MODEL, args=torch.randn(1, 3, 64, 64), f='output/checkpoints/model.onnx', verbose=True, input_names=['input'], output_names=['output'])
+torch.onnx.export(
+    model=MODEL,
+    args=torch.randn(1, 3, 64, 64),
+    f="output/checkpoints/model.onnx",
+    verbose=True,
+    input_names=["input"],
+    output_names=["output"],
+)
 
-onnx2tf.convert(input_onnx_file_path='output/checkpoints/model.onnx', output_folder_path='output/checkpoints/converted/')
+onnx2tf.convert(
+    input_onnx_file_path="output/checkpoints/model.onnx",
+    output_folder_path="output/checkpoints/converted/",
+)

eval.py

@@ -18,6 +18,7 @@ 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
@@ -74,5 +75,6 @@ def predict_image(image_path, model, transform):
     conf_matrix.plot()
     plt.show()
 
+
 # Call predict_image function
 predict_image(image_path, MODEL, preprocess)

models.py

@@ -39,4 +39,3 @@ 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

@@ -15,7 +15,10 @@ def setup_tensorboard():
 
 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
+        RAW_DATA_DIR + str(TASK),
+        AUG_DATA_DIR + str(TASK),
+        EXTERNAL_DATA_DIR + str(TASK),
+        preprocess,
     )
 
 
@@ -194,4 +197,4 @@ def main_training_loop():
 
 
 if __name__ == "__main__":
-    main_training_loop()
+    main_training_loop()

tuning.py

@@ -115,9 +115,7 @@ if __name__ == "__main__":
     )
 
     # Optimize the hyperparameters
-    study.optimize(
-        objective, n_trials=N_TRIALS, timeout=TIMEOUT
-    )
+    study.optimize(objective, n_trials=N_TRIALS, timeout=TIMEOUT)
 
     # Print the best trial
     best_trial = study.best_trial