FEAT: Add Type Handler on Function and Clas

main.py

@@ -11,7 +11,7 @@ import torch
 import torch.nn as nn
 import matplotlib.pyplot as plt
 from torchvision import transforms, datasets
-from torch.utils.data import DataLoader, random_split
+from torch.utils.data import DataLoader, random_split, Dataset
 
 # Import custom modules
 sys.path.append("./utils")
@@ -23,7 +23,7 @@ from Trainer import model_train
 
 
 # Set random seeds for reproducibility
-def set_seed(seed=42):
+def set_seed(seed=42) -> None:
     """Set seeds for reproducibility."""
     random.seed(seed)
     np.random.seed(seed)
@@ -35,7 +35,7 @@ def set_seed(seed=42):
         torch.backends.cudnn.benchmark = False
 
 
-def get_transform():
+def get_transform() -> transforms.Compose:
     """
     Get standard data transform for both training and validation/testing.
 
@@ -55,7 +55,7 @@ def get_transform():
     return transform
 
 
-def load_data(args):
+def load_data(args) -> tuple:
     """
     Load and prepare datasets from separate directories for training and evaluation.
 
@@ -135,7 +135,13 @@ def load_data(args):
     return train_loader, val_loader, test_loader, eval_dataset
 
 
-def train_single_model(args, train_loader, val_loader, test_loader, dataset):
+def train_single_model(
+    args,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    test_loader: DataLoader,
+    dataset: Dataset,
+) -> None:
     """Train a single model specified by the arguments."""
 
     print(f"Creating {args.model} model...")
@@ -177,7 +183,13 @@ def train_single_model(args, train_loader, val_loader, test_loader, dataset):
         print("Training failed. Cannot evaluate on test set.")
 
 
-def compare_multiple_models(args, train_loader, val_loader, test_loader, dataset):
+def compare_multiple_models(
+    args,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    test_loader: DataLoader,
+    dataset: Dataset,
+) -> None:
     """Compare multiple models."""
 
     print("Preparing to compare multiple models...")
@@ -217,7 +229,7 @@ def compare_multiple_models(args, train_loader, val_loader, test_loader, dataset
     )
 
 
-def main():
+def main() -> None:
     """Main function to run the eye disease detection application."""
 
     # Set up argument parser with example usage

utils/Comparator.py

@@ -1,10 +1,28 @@
+from torch.utils.data import DataLoader, Dataset
 from Trainer import model_train
 from Evaluator import ClassificationEvaluator
 
 
 def compare_models(
-    models, train_loader, val_loader, test_loader, dataset, epochs=20, names=None
-):
+    models: list,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    test_loader: DataLoader,
+    dataset: Dataset,
+    epochs: int = 20,
+    names: list = None,
+) -> None:
+    """
+    Compare multiple models on validation and test datasets.
+    Args:
+        models (list): List of models to compare.
+        train_loader (DataLoader): DataLoader for training data.
+        val_loader (DataLoader): DataLoader for validation data.
+        test_loader (DataLoader): DataLoader for test data.
+        dataset (Dataset): Dataset object containing class names.
+        epochs (int): Number of epochs for training.
+        names (list): List of model names. If None, default names will be used.
+    """
     if names is None:
         names = [f"Model {i+1}" for i in range(len(models))]
 

utils/Evaluator.py

@@ -1,4 +1,6 @@
 import numpy as np
+import torch.nn as nn
+from torch.utils.data import DataLoader
 import torch
 import matplotlib.pyplot as plt
 import seaborn as sns
@@ -23,7 +25,7 @@ class ClassificationEvaluator:
     and generate visualizations for model evaluation.
     """
 
-    def __init__(self, class_names):
+    def __init__(self, class_names: list):
         """
         Initialize the evaluator with class names.
 
@@ -39,7 +41,7 @@ class ClassificationEvaluator:
             return data.cpu().numpy()
         return np.array(data)
 
-    def evaluate_model(self, model, test_loader):
+    def evaluate_model(self, model: nn.Module, test_loader: DataLoader) -> dict:
         """
         Evaluate a trained model on test dataset.
 
@@ -75,7 +77,7 @@ class ClassificationEvaluator:
         results = self.compute_metrics(all_labels, all_preds, all_scores)
         return results
 
-    def compute_metrics(self, y_true, y_pred, y_scores, model_name=""):
+    def compute_metrics(self, y_true, y_pred, y_scores, model_name: str = "") -> dict:
         """
         Compute comprehensive classification metrics.
 
@@ -135,7 +137,7 @@ class ClassificationEvaluator:
             "kappa": kappa,
         }
 
-    def plot_roc_curves(self, y_true, y_scores):
+    def plot_roc_curves(self, y_true, y_scores) -> dict:
         """
         Plot ROC curves for multi-class classification.
 
@@ -197,7 +199,7 @@ class ClassificationEvaluator:
 
         return roc_auc
 
-    def plot_pr_curves(self, y_true, y_scores):
+    def plot_pr_curves(self, y_true, y_scores) -> dict:
         """
         Plot Precision-Recall curves for multi-class classification.
 
@@ -263,7 +265,7 @@ class ClassificationEvaluator:
 
         return avg_precision
 
-    def plot_confusion_matrix(self, y_true, y_pred):
+    def plot_confusion_matrix(self, y_true, y_pred) -> None:
         """
         Plot confusion matrix.
 
@@ -296,7 +298,7 @@ class ClassificationEvaluator:
         plt.tight_layout()
         plt.show()
 
-    def plot_per_class_accuracy(self, y_true, y_pred):
+    def plot_per_class_accuracy(self, y_true, y_pred) -> np.ndarray:
         """
         Plot per-class accuracy.
 
@@ -328,7 +330,9 @@ class ClassificationEvaluator:
 
         return per_class_accuracy
 
-    def plot_training_history(self, train_losses, val_losses, train_accs, val_accs):
+    def plot_training_history(
+        self, train_losses, val_losses, train_accs, val_accs
+    ) -> None:
         """
         Plot accuracy and loss curves from training history.
 

utils/ModelCreator.py

@@ -1,13 +1,23 @@
 import torch
+from torch import device
 import torch.nn as nn
 import timm
 
 # Set device
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+DEVICE = device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-class EyeDetectionModels:
-    def __init__(self, num_classes, freeze_layers=True, device=DEVICE):
+class EyeDetectionModels(object):
+    """
+    A class to create and configure various deep learning models for eye detection tasks.
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        freeze_layers: bool = True,
+        device: device = DEVICE,
+    ):
         """
         Initialize the EyeDetectionModels class.
         This class provides methods to create and configure various deep learning models for eye detection.
@@ -26,7 +36,7 @@ class EyeDetectionModels:
 
     # Model architecture functions
     @staticmethod
-    def _get_feature_blocks(model):
+    def _get_feature_blocks(model: nn.Module) -> nn.ModuleList:
         """
         Utility: locate the main feature blocks container in a timm model.
         Returns a list-like module of blocks.
@@ -38,14 +48,14 @@ class EyeDetectionModels:
         return list(model.children())[:-1]
 
     @staticmethod
-    def _freeze_except_last_n(blocks, n):
+    def _freeze_except_last_n(blocks: nn.ModuleList, n: int) -> None:
         total = len(blocks)
         for idx, block in enumerate(blocks):
             requires = idx >= total - n
             for p in block.parameters():
                 p.requires_grad = requires
 
-    def get_model_mobilenetv4(self):
+    def get_model_mobilenetv4(self) -> nn.Module:
         model = timm.create_model(
             "mobilenetv4_conv_medium.e500_r256_in1k", pretrained=True
         )
@@ -62,14 +72,12 @@ class EyeDetectionModels:
         )
         return model.to(self.device)
 
-    def get_model_levit(self):
+    def get_model_levit(self) -> nn.Module:
         model = timm.create_model("levit_128s.fb_dist_in1k", pretrained=True)
         if self.freeze_layers:
             blocks = self._get_feature_blocks(model)
             self._freeze_except_last_n(blocks, 2)
         # Attempt to extract in_features from model.head or classifier
-        head = getattr(model, "head_dist", None) or getattr(model, "classifier", None)
-        linear = getattr(head, "linear")
         in_features = 384
         model.head = nn.Sequential(
             nn.Linear(in_features, 512),
@@ -85,15 +93,12 @@ class EyeDetectionModels:
         )
         return model.to(self.device)
 
-    def get_model_efficientvit(self):
+    def get_model_efficientvit(self) -> nn.Module:
         model = timm.create_model("efficientvit_m1.r224_in1k", pretrained=True)
         if self.freeze_layers:
             blocks = self._get_feature_blocks(model)
             self._freeze_except_last_n(blocks, 2)
         # handle different head naming
-        head = getattr(model, "head", None)
-        print(head)
-        linear = getattr(head, "linear")
         in_features = 192
         model.head.linear = nn.Sequential(
             nn.Linear(in_features, 512),
@@ -103,7 +108,7 @@ class EyeDetectionModels:
         )
         return model.to(self.device)
 
-    def get_model_gernet(self):
+    def get_model_gernet(self) -> nn.Module:
         """
         Load and configure a GENet (General and Efficient Network) model with customizable classifier.
 
@@ -142,7 +147,7 @@ class EyeDetectionModels:
         )
         return model.to(self.device)
 
-    def get_model_regnetx(self):
+    def get_model_regnetx(self) -> nn.Module:
         """
         Load and configure a RegNetX model with customizable classifier.
 

utils/Trainer.py

@@ -3,6 +3,7 @@ import torch
 import torch.nn.functional as F
 import torch.optim as optim
 import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
 from tqdm import tqdm
 import gc
 
@@ -11,16 +12,16 @@ from Callback import EarlyStopping
 
 
 def train_model(
-    model,
-    criterion,
-    optimizer,
-    scheduler,
-    train_loader,
-    val_loader,
-    early_stopping,
-    epochs=15,
-    use_ddp=False,
-):
+    model: nn.Module,
+    criterion: nn.Module,
+    optimizer: optim.Optimizer,
+    scheduler: optim.lr_scheduler._LRScheduler,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    early_stopping: EarlyStopping,
+    epochs: int = 15,
+    use_ddp: bool = False,
+) -> tuple:
     """
     Train the model and perform validation using multiple GPUs.
     Supports both DataParallel (DP) and DistributedDataParallel (DDP) modes.
@@ -180,7 +181,13 @@ def train_model(
     )
 
 
-def model_train(model, train_loader, val_loader, dataset, epochs=20):
+def model_train(
+    model: nn.Module,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    dataset: Dataset,
+    epochs: int = 20,
+) -> tuple:
     model_name = type(model).__name__
     if hasattr(model, "pretrained_cfg") and "name" in model.pretrained_cfg:
         model_name = model.pretrained_cfg["name"]