Upload 29 files

Models_Conversion/base_fusion_model.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from transformers import ConvNextModel
+from huggingface_hub import hf_hub_download
+
+
+# ==========================================
+# 1. MODEL ARCHITECTURE DEFINITION
+# ==========================================
+class FusionClassifier(nn.Module):
+    def __init__(self, num_classes, convnext_model_name="facebook/convnext-small-224"):
+        super().__init__()
+
+        # EfficientNet-V2-S backbone
+        eff = models.efficientnet_v2_s(weights=models.EfficientNet_V2_S_Weights.IMAGENET1K_V1)
+
+        for param in eff.parameters():
+            param.requires_grad = False
+
+        for param in eff.features[5].parameters():
+            param.requires_grad = True
+        for param in eff.features[6].parameters():
+            param.requires_grad = True
+        for param in eff.features[7].parameters():
+            param.requires_grad = True
+
+        self.eff_features = eff.features
+        self.eff_avgpool = eff.avgpool
+        self.eff_out_dim = eff.classifier[1].in_features
+
+        # ConvNeXt-Small backbone
+        cnx = ConvNextModel.from_pretrained(convnext_model_name)
+
+        for param in cnx.parameters():
+            param.requires_grad = False
+
+        for param in cnx.encoder.stages[2].parameters():
+            param.requires_grad = True
+        for param in cnx.encoder.stages[3].parameters():
+            param.requires_grad = True
+        for param in cnx.layernorm.parameters():
+            param.requires_grad = True
+
+        self.cnx_backbone = cnx
+        self.cnx_out_dim = 768
+
+        # Fusion head
+        fused_dim = self.eff_out_dim + self.cnx_out_dim
+        self.fusion_head = nn.Sequential(
+            nn.Dropout(p=0.4),
+            nn.Linear(fused_dim, 512),
+            nn.LayerNorm(512),
+            nn.GELU(),
+            nn.Dropout(p=0.3),
+            nn.Linear(512, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(p=0.2),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, pixel_values_eff, pixel_values_cnx):
+        x_eff = self.eff_features(pixel_values_eff)
+        x_eff = self.eff_avgpool(x_eff)
+        x_eff = torch.flatten(x_eff, 1)
+
+        cnx_out = self.cnx_backbone(pixel_values=pixel_values_cnx, return_dict=True)
+        x_cnx = cnx_out.pooler_output
+
+        fused = torch.cat([x_eff, x_cnx], dim=1)
+        logits = self.fusion_head(fused)
+
+        return logits
+
+
+# ==========================================
+# 2. HUGGING FACE DOWNLOAD & LOAD FUNCTION
+# ==========================================
+def load_fusion_model_from_hf(
+    repo_id: str = "junaid17/best_fusion_model_fp16",
+    filename: str = "best_fusion_model_fp16.pt",
+    num_classes: int = 10,  # Replace with your actual number of classes
+    device: str = "cpu",
+) -> nn.Module:
+    """Downloads weights from Hugging Face Hub and loads into FusionClassifier."""
+    print(f"Downloading checkpoint from Hugging Face Hub: '{repo_id}/{filename}'...")
+    
+    # Download weights file from Hugging Face Hub
+    checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    
+    # Instantiate Model
+    model = FusionClassifier(num_classes=num_classes)
+    
+    # Load weights
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+        model.load_state_dict(checkpoint["state_dict"])
+    elif isinstance(checkpoint, dict):
+        model.load_state_dict(checkpoint)
+    else:
+        model = checkpoint
+
+    # Ensure model is float32 for stable ONNX export
+    model = model.float().to(device)
+    model.eval()
+    
+    print("✅ Model loaded successfully from Hugging Face.")
+    return model

Models_Conversion/base_model_resnet.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+from huggingface_hub import hf_hub_download
+
+
+class Car_Classifier_Resnet(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+
+        self.model = models.resnet18(weights="DEFAULT")
+
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        for param in self.model.layer3.parameters():
+            param.requires_grad = True
+
+        for param in self.model.layer4.parameters():
+            param.requires_grad = True
+
+        self.model.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(self.model.fc.in_features, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+def load_resnet_model_from_hf(
+    repo_id: str,
+    filename: str = "car-damage-classifier.pt",
+    num_classes: int = 6,
+    device: str = "cpu",
+    hf_token: str = None,
+) -> torch.nn.Module:
+    """
+    Downloads and loads the state_dict for Car_Classifier_Resnet from Hugging Face Hub.
+    """
+    print(f"Downloading checkpoint '{filename}' from Hugging Face repo '{repo_id}'...")
+    checkpoint_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        token=hf_token
+    )
+
+    model = Car_Classifier_Resnet(num_classes=num_classes)
+    state_dict = torch.load(checkpoint_path, map_location=device)
+
+    if isinstance(state_dict, dict) and "state_dict" in state_dict:
+        state_dict = state_dict["state_dict"]
+
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    print("✅ ResNet-18 model loaded successfully.")
+    return model

Models_Conversion/base_model_yolo.py

@@ -0,0 +1,22 @@
+import os
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
+
+
+def load_yolo_model_from_hf(
+    repo_id: str = "junaid17/Yolo_Model",
+    filename: str = "damage_detector.pt",
+    hf_token: str = None,
+) -> YOLO:
+    """
+    Downloads and loads a trained Ultralytics YOLO model from Hugging Face Hub.
+    """
+    print(f"Downloading YOLO checkpoint '{filename}' from Hugging Face repo '{repo_id}'...")
+    checkpoint_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        token=hf_token
+    )
+    model = YOLO(checkpoint_path)
+    print("✅ YOLO model loaded successfully.")
+    return model

Models_Conversion/base_vs_onnx_model_evaluation.ipynb

@@ -0,0 +1,532 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "7c5a22e0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import onnxruntime as ort\n",
+    "from tqdm import tqdm\n",
+    "from torch.utils.data import DataLoader, Dataset\n",
+    "from torchvision import transforms\n",
+    "from torchvision.datasets import ImageFolder\n",
+    "from transformers import ConvNextImageProcessor\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report\n",
+    "from base_fusion_model import FusionClassifier\n",
+    "from base_model_resnet import Car_Classifier_Resnet\n",
+    "\n",
+    "# ==========================================\n",
+    "# IMPORT YOUR BASE PYTORCH MODELS HERE\n",
+    "# Example: from my_models import FusionClassifier, Car_Classifier_Resnet\n",
+    "# ==========================================\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "HF_TOKEN = os.getenv(\"HF_TOKEN\")  # Ensure your Hugging Face token is set in environment\n",
+    "\n",
+    "# ==========================================\n",
+    "# 1. IMAGE PROCESSING & TRANSFORMS\n",
+    "# ==========================================\n",
+    "\n",
+    "# --- Fusion Model Transforms ---\n",
+    "IMG_SIZE = 260\n",
+    "convnext_model_name = \"facebook/convnext-small-224\"\n",
+    "processor = ConvNextImageProcessor.from_pretrained(convnext_model_name)\n",
+    "\n",
+    "fusion_val_transforms = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "])\n",
+    "\n",
+    "eff_normalize = transforms.Compose([\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.485, 0.456, 0.406],\n",
+    "                         [0.229, 0.224, 0.225])\n",
+    "])\n",
+    "\n",
+    "# --- ResNet Model Transforms ---\n",
+    "resnet_test_transforms = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.485, 0.456, 0.406],\n",
+    "                         [0.229, 0.224, 0.225])\n",
+    "])\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 2. DATASETS\n",
+    "# ==========================================\n",
+    "\n",
+    "class FusionDataset(Dataset):\n",
+    "    def __init__(self, data_dir, processor, transform=None):\n",
+    "        self.dataset = ImageFolder(root=data_dir, transform=transform)\n",
+    "        self.processor = processor\n",
+    "        self.classes = self.dataset.classes\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        pil_image, label = self.dataset[idx]\n",
+    "\n",
+    "        # EfficientNet branch\n",
+    "        pixel_eff = eff_normalize(pil_image)\n",
+    "\n",
+    "        # ConvNeXt branch\n",
+    "        enc = self.processor(images=pil_image, return_tensors='pt')\n",
+    "        pixel_cnx = enc['pixel_values'].squeeze(0)\n",
+    "\n",
+    "        return {\n",
+    "            'pixel_values_eff': pixel_eff,\n",
+    "            'pixel_values_cnx': pixel_cnx,\n",
+    "            'labels': torch.tensor(label)\n",
+    "        }\n",
+    "\n",
+    "class ResNetDataset(Dataset):\n",
+    "    def __init__(self, data_dir, transform=None):\n",
+    "        self.dataset = ImageFolder(root=data_dir, transform=transform)\n",
+    "        self.classes = self.dataset.classes\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        return self.dataset[idx]\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 3. HUGGING FACE MODEL LOADER\n",
+    "# ==========================================\n",
+    "def load_models_from_hf(repo_id, pt_filename, onnx_filename, model_instance):\n",
+    "    \"\"\"\n",
+    "    Downloads PyTorch and ONNX models from HF and loads PyTorch weights into the base instance.\n",
+    "    \"\"\"\n",
+    "    print(f\"\\n⬇️ Downloading models from Hugging Face: {repo_id}\")\n",
+    "    \n",
+    "    # 1. Download and load PyTorch weights\n",
+    "    pt_path = hf_hub_download(repo_id=repo_id, filename=pt_filename, token=HF_TOKEN)\n",
+    "    checkpoint = torch.load(pt_path, map_location=device)\n",
+    "    \n",
+    "    # Handle both full checkpoints dicts and direct state_dicts\n",
+    "    if 'model_state_dict' in checkpoint:\n",
+    "        model_instance.load_state_dict(checkpoint['model_state_dict'])\n",
+    "        print(f\"✅ PyTorch model loaded (Val Acc: {checkpoint.get('val_acc', 'N/A')}%)\")\n",
+    "    else:\n",
+    "        model_instance.load_state_dict(checkpoint)\n",
+    "        print(\"✅ PyTorch model loaded\")\n",
+    "        \n",
+    "    model_instance.to(device)\n",
+    "    model_instance.eval()\n",
+    "\n",
+    "    # 2. Download ONNX model\n",
+    "    onnx_path = hf_hub_download(repo_id=repo_id, filename=onnx_filename, token=HF_TOKEN)\n",
+    "    print(\"✅ ONNX model downloaded\")\n",
+    "    \n",
+    "    return model_instance, onnx_path\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 4. EVALUATION ENGINE\n",
+    "# ==========================================\n",
+    "def evaluate_models(model_type, torch_model, onnx_model_path, dataloader, target_names):\n",
+    "    \"\"\"\n",
+    "    Runs evaluation for both PyTorch and ONNX models side-by-side.\n",
+    "    model_type must be either 'fusion' or 'resnet'.\n",
+    "    \"\"\"\n",
+    "    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']\n",
+    "    ort_session = ort.InferenceSession(onnx_model_path, providers=providers)\n",
+    "    \n",
+    "    all_labels = []\n",
+    "    torch_preds = []\n",
+    "    onnx_preds = []\n",
+    "\n",
+    "    print(f\"\\n🚀 Starting Evaluation for {model_type.upper()} Model...\")\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        for batch in tqdm(dataloader, desc=\"Evaluating\"):\n",
+    "            if model_type == 'fusion':\n",
+    "                images_eff = batch['pixel_values_eff'].to(device)\n",
+    "                images_cnx = batch['pixel_values_cnx'].to(device)\n",
+    "                labels = batch['labels'].numpy()\n",
+    "                \n",
+    "                # PyTorch Inference\n",
+    "                logits_torch = torch_model(images_eff, images_cnx)\n",
+    "                \n",
+    "                # ONNX Inference (Requires inputs matching the exported graph)\n",
+    "                onnx_inputs = {\n",
+    "                    ort_session.get_inputs()[0].name: images_eff.cpu().numpy(),\n",
+    "                    ort_session.get_inputs()[1].name: images_cnx.cpu().numpy()\n",
+    "                }\n",
+    "                logits_onnx = ort_session.run(None, onnx_inputs)[0]\n",
+    "\n",
+    "            elif model_type == 'resnet':\n",
+    "                images, labels_tensor = batch\n",
+    "                images = images.to(device)\n",
+    "                labels = labels_tensor.numpy()\n",
+    "                \n",
+    "                # PyTorch Inference\n",
+    "                logits_torch = torch_model(images)\n",
+    "                \n",
+    "                # ONNX Inference\n",
+    "                onnx_inputs = {ort_session.get_inputs()[0].name: images.cpu().numpy()}\n",
+    "                logits_onnx = ort_session.run(None, onnx_inputs)[0]\n",
+    "\n",
+    "            # Collect Predictions\n",
+    "            preds_t = torch.argmax(logits_torch, dim=1).cpu().numpy()\n",
+    "            preds_o = np.argmax(logits_onnx, axis=1)\n",
+    "\n",
+    "            all_labels.extend(labels)\n",
+    "            torch_preds.extend(preds_t)\n",
+    "            onnx_preds.extend(preds_o)\n",
+    "\n",
+    "    # Helper function to print detailed metrics\n",
+    "    def print_metrics(y_true, y_pred, title):\n",
+    "        acc = accuracy_score(y_true, y_pred)\n",
+    "        prec = precision_score(y_true, y_pred, average='macro', zero_division=0)\n",
+    "        rec = recall_score(y_true, y_pred, average='macro', zero_division=0)\n",
+    "        f1 = f1_score(y_true, y_pred, average='macro', zero_division=0)\n",
+    "\n",
+    "        print(f\"\\n{'='*50}\")\n",
+    "        print(f\" {title.center(48)} \")\n",
+    "        print(f\"{'='*50}\")\n",
+    "        print(f\"Accuracy  : {acc:.4f}\")\n",
+    "        print(f\"Precision : {prec:.4f}\")\n",
+    "        print(f\"Recall    : {rec:.4f}\")\n",
+    "        print(f\"F1 Score  : {f1:.4f}\")\n",
+    "        print(\"-\" * 50)\n",
+    "        print(\"Classification Report:\")\n",
+    "        print(classification_report(y_true, y_pred, target_names=target_names, zero_division=0))\n",
+    "\n",
+    "    print_metrics(all_labels, torch_preds, f\"{model_type.upper()} - PYTORCH MODEL\")\n",
+    "    print_metrics(all_labels, onnx_preds, f\"{model_type.upper()} - ONNX MODEL\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "375a4992",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c428a2fa573042d590a7663a583b9c97",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading weights:   0%|          | 0/342 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[transformers] \u001b[1mConvNextModel LOAD REPORT\u001b[0m from: facebook/convnext-small-224\n",
+      "Key               | Status     |  | \n",
+      "------------------+------------+--+-\n",
+      "classifier.bias   | UNEXPECTED |  | \n",
+      "classifier.weight | UNEXPECTED |  | \n",
+      "\n",
+      "Notes:\n",
+      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "⬇️ Downloading models from Hugging Face: junaid17/best_fusion_model_fp16\n",
+      "✅ PyTorch model loaded\n",
+      "✅ ONNX model downloaded\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\DamageLensAI\\myvenv\\Lib\\site-packages\\onnxruntime\\capi\\onnxruntime_inference_collection.py:153: UserWarning: Specified provider 'CUDAExecutionProvider' is not in available provider names.Available providers: 'AzureExecutionProvider, CPUExecutionProvider'\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "🚀 Starting Evaluation for FUSION Model...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Evaluating: 100%|██████████| 29/29 [04:30<00:00,  9.31s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "==================================================\n",
+      "              FUSION - PYTORCH MODEL              \n",
+      "==================================================\n",
+      "Accuracy  : 0.8413\n",
+      "Precision : 0.8364\n",
+      "Recall    : 0.8342\n",
+      "F1 Score  : 0.8348\n",
+      "--------------------------------------------------\n",
+      "Classification Report:\n",
+      "              precision    recall  f1-score   support\n",
+      "\n",
+      "  F_Breakage       0.86      0.86      0.86       100\n",
+      "   F_Crushed       0.78      0.78      0.78        80\n",
+      "    F_Normal       0.91      0.92      0.92       100\n",
+      "  R_Breakage       0.83      0.80      0.81        60\n",
+      "   R_Crushed       0.74      0.82      0.78        60\n",
+      "    R_Normal       0.89      0.83      0.86        60\n",
+      "\n",
+      "    accuracy                           0.84       460\n",
+      "   macro avg       0.84      0.83      0.83       460\n",
+      "weighted avg       0.84      0.84      0.84       460\n",
+      "\n",
+      "\n",
+      "==================================================\n",
+      "               FUSION - ONNX MODEL                \n",
+      "==================================================\n",
+      "Accuracy  : 0.8413\n",
+      "Precision : 0.8364\n",
+      "Recall    : 0.8342\n",
+      "F1 Score  : 0.8348\n",
+      "--------------------------------------------------\n",
+      "Classification Report:\n",
+      "              precision    recall  f1-score   support\n",
+      "\n",
+      "  F_Breakage       0.86      0.86      0.86       100\n",
+      "   F_Crushed       0.78      0.78      0.78        80\n",
+      "    F_Normal       0.91      0.92      0.92       100\n",
+      "  R_Breakage       0.83      0.80      0.81        60\n",
+      "   R_Crushed       0.74      0.82      0.78        60\n",
+      "    R_Normal       0.89      0.83      0.86        60\n",
+      "\n",
+      "    accuracy                           0.84       460\n",
+      "   macro avg       0.84      0.83      0.83       460\n",
+      "weighted avg       0.84      0.84      0.84       460\n",
+      "\n",
+      "\n",
+      "⬇️ Downloading models from Hugging Face: junaid17/car-damage-classifier\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "477e596952ee4da99a21c89acb64e70e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "car-damage-classifier.pt: reconstructing file:   0%|          |  0.00B / 45.3MB            "
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "eb1e5cde28db4feb9c1c4fba02eaef64",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "car-damage-classifier.pt: downloading bytes:           |  0.00B            "
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\DamageLensAI\\myvenv\\Lib\\site-packages\\huggingface_hub\\file_download.py:141: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--car-damage-classifier. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ PyTorch model loaded\n",
+      "✅ ONNX model downloaded\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\DamageLensAI\\myvenv\\Lib\\site-packages\\onnxruntime\\capi\\onnxruntime_inference_collection.py:153: UserWarning: Specified provider 'CUDAExecutionProvider' is not in available provider names.Available providers: 'AzureExecutionProvider, CPUExecutionProvider'\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "🚀 Starting Evaluation for RESNET Model...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Evaluating: 100%|██████████| 29/29 [00:35<00:00,  1.21s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "==================================================\n",
+      "              RESNET - PYTORCH MODEL              \n",
+      "==================================================\n",
+      "Accuracy  : 0.7739\n",
+      "Precision : 0.7734\n",
+      "Recall    : 0.7661\n",
+      "F1 Score  : 0.7681\n",
+      "--------------------------------------------------\n",
+      "Classification Report:\n",
+      "              precision    recall  f1-score   support\n",
+      "\n",
+      "  F_Breakage       0.83      0.78      0.80       100\n",
+      "   F_Crushed       0.63      0.75      0.69        80\n",
+      "    F_Normal       0.89      0.85      0.87       100\n",
+      "  R_Breakage       0.81      0.72      0.76        60\n",
+      "   R_Crushed       0.67      0.72      0.69        60\n",
+      "    R_Normal       0.81      0.78      0.80        60\n",
+      "\n",
+      "    accuracy                           0.77       460\n",
+      "   macro avg       0.77      0.77      0.77       460\n",
+      "weighted avg       0.78      0.77      0.78       460\n",
+      "\n",
+      "\n",
+      "==================================================\n",
+      "               RESNET - ONNX MODEL                \n",
+      "==================================================\n",
+      "Accuracy  : 0.7739\n",
+      "Precision : 0.7734\n",
+      "Recall    : 0.7661\n",
+      "F1 Score  : 0.7681\n",
+      "--------------------------------------------------\n",
+      "Classification Report:\n",
+      "              precision    recall  f1-score   support\n",
+      "\n",
+      "  F_Breakage       0.83      0.78      0.80       100\n",
+      "   F_Crushed       0.63      0.75      0.69        80\n",
+      "    F_Normal       0.89      0.85      0.87       100\n",
+      "  R_Breakage       0.81      0.72      0.76        60\n",
+      "   R_Crushed       0.67      0.72      0.69        60\n",
+      "    R_Normal       0.81      0.78      0.80        60\n",
+      "\n",
+      "    accuracy                           0.77       460\n",
+      "   macro avg       0.77      0.77      0.77       460\n",
+      "weighted avg       0.78      0.77      0.78       460\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "if __name__ == \"__main__\":\n",
+    "    eval_dir = \"D:\\\\DamageLensAI\\\\val\"  # Adjust to your evaluation folder\n",
+    "\n",
+    "    # ---------------------------------------------------------\n",
+    "    # EVALUATE FUSION MODEL\n",
+    "    # ---------------------------------------------------------\n",
+    "    fusion_dataset = FusionDataset(data_dir=eval_dir, processor=processor, transform=fusion_val_transforms)\n",
+    "    fusion_loader = DataLoader(fusion_dataset, batch_size=16, shuffle=False)\n",
+    "    \n",
+    "    # Initialize your base Fusion PyTorch architecture\n",
+    "    # Replace 'FusionClassifier' with your actual class name and arguments\n",
+    "    base_fusion = FusionClassifier(num_classes=len(fusion_dataset.classes), convnext_model_name=convnext_model_name)\n",
+    "    \n",
+    "    # Load weights and get ONNX path from Hub\n",
+    "    torch_fusion, onnx_fusion_path = load_models_from_hf(\n",
+    "        repo_id=\"junaid17/best_fusion_model_fp16\",\n",
+    "        pt_filename=\"best_fusion_model_fp16.pt\", # Update if named differently on HF\n",
+    "        onnx_filename=\"fusion_model.onnx\",\n",
+    "        model_instance=base_fusion\n",
+    "    )\n",
+    "    \n",
+    "    evaluate_models('fusion', torch_fusion, onnx_fusion_path, fusion_loader, fusion_dataset.classes)\n",
+    "\n",
+    "    # ---------------------------------------------------------\n",
+    "    # EVALUATE RESNET MODEL\n",
+    "    # ---------------------------------------------------------\n",
+    "    resnet_dataset = ResNetDataset(data_dir=eval_dir, transform=resnet_test_transforms)\n",
+    "    resnet_loader = DataLoader(resnet_dataset, batch_size=16, shuffle=False)\n",
+    "    \n",
+    "    # Initialize your base ResNet PyTorch architecture\n",
+    "    # Replace 'get_resnet_model' with how you initialize your ResNet\n",
+    "    base_resnet = Car_Classifier_Resnet(num_classes=len(resnet_dataset.classes))\n",
+    "    \n",
+    "    # Load weights and get ONNX path from Hub\n",
+    "    torch_resnet, onnx_resnet_path = load_models_from_hf(\n",
+    "        repo_id=\"junaid17/car-damage-classifier\",\n",
+    "        pt_filename=\"car-damage-classifier.pt\", # Update if named differently on HF\n",
+    "        onnx_filename=\"car-damage-classifier.onnx\",\n",
+    "        model_instance=base_resnet\n",
+    "    )\n",
+    "    \n",
+    "    evaluate_models('resnet', torch_resnet, onnx_resnet_path, resnet_loader, resnet_dataset.classes)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36697686",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "myvenv (3.11.0)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Models_Conversion/model_onnx_conversion.ipynb

@@ -0,0 +1,1171 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ac6e10d4",
+   "metadata": {},
+   "source": [
+    "# Fusion Model Conversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "0f63f311",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading checkpoint from Hugging Face Hub: 'junaid17/best_fusion_model_fp16/best_fusion_model_fp16.pt'...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading weights: 100%|██████████| 342/342 [00:00<00:00, 1952.82it/s]\n",
+      "[transformers] \u001b[1mConvNextModel LOAD REPORT\u001b[0m from: facebook/convnext-small-224\n",
+      "Key               | Status     |  | \n",
+      "------------------+------------+--+-\n",
+      "classifier.weight | UNEXPECTED |  | \n",
+      "classifier.bias   | UNEXPECTED |  | \n",
+      "\n",
+      "Notes:\n",
+      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Model loaded successfully from Hugging Face.\n",
+      "Exporting Fusion Model to temporary ONNX graph...\n",
+      "✅ Base ONNX export successful.\n",
+      "Compressing ONNX model to FP16 precision...\n",
+      "✅ ONNX graph compressed to FP16 successfully.\n",
+      "Uploading 'fusion_model.onnx' to Hugging Face repo: 'junaid17/best_fusion_model_fp16'...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Processing Files (1 / 1): 100%|██████████|  142MB /  142MB,  680kB/s  \n",
+      "New Data Upload: 100%|██████████|  142MB /  142MB,  680kB/s  \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🚀 Successfully uploaded FP16 model to https://huggingface.co/junaid17/best_fusion_model_fp16/blob/main/fusion_model.onnx\n",
+      "Downloading 'fusion_model.onnx' from Hugging Face for evaluation...\n",
+      "📸 Found 60 images in 'D:\\DamageLensAI\\test_images'. Starting benchmark...\n",
+      "\n",
+      "===========================================================================\n",
+      "IMAGE NAME                | PYTORCH (ms)   | ONNX (ms)      | MATCH?    \n",
+      "===========================================================================\n",
+      "FB_10.jpg                 | 775.56         | 225.80         | ✅ YES     \n",
+      "FB_11.jpg                 | 683.06         | 263.21         | ✅ YES     \n",
+      "FB_13.jpg                 | 359.21         | 156.10         | ✅ YES     \n",
+      "FB_15.jpg                 | 357.64         | 152.04         | ✅ YES     \n",
+      "FB_16.jpg                 | 340.61         | 157.13         | ✅ YES     \n",
+      "FB_19.jpg                 | 343.34         | 157.87         | ✅ YES     \n",
+      "FB_20.jpg                 | 562.84         | 188.25         | ✅ YES     \n",
+      "FB_4.jpg                  | 328.80         | 162.34         | ✅ YES     \n",
+      "FB_5.jpg                  | 466.30         | 166.86         | ✅ YES     \n",
+      "FB_9.jpg                  | 657.23         | 188.45         | ✅ YES     \n",
+      "FC_1.jpg                  | 372.87         | 234.55         | ✅ YES     \n",
+      "FC_10.jpg                 | 473.04         | 164.85         | ✅ YES     \n",
+      "FC_11.jpg                 | 332.92         | 154.70         | ✅ YES     \n",
+      "FC_12.jpg                 | 379.71         | 197.05         | ✅ YES     \n",
+      "FC_3.jpg                  | 546.80         | 297.15         | ✅ YES     \n",
+      "FC_4.jpg                  | 810.00         | 387.67         | ✅ YES     \n",
+      "FC_6.jpg                  | 528.11         | 211.10         | ✅ YES     \n",
+      "FC_7.jpg                  | 546.06         | 267.98         | ✅ YES     \n",
+      "FC_8.jpg                  | 665.58         | 230.13         | ✅ YES     \n",
+      "FC_9.jpg                  | 516.99         | 213.58         | ✅ YES     \n",
+      "FN_1.jpg                  | 497.13         | 216.17         | ✅ YES     \n",
+      "FN_10.jpg                 | 484.10         | 215.26         | ✅ YES     \n",
+      "FN_11.jpg                 | 521.18         | 221.95         | ✅ YES     \n",
+      "FN_12.jpg                 | 498.07         | 235.80         | ✅ YES     \n",
+      "FN_3.jpg                  | 473.85         | 219.38         | ✅ YES     \n",
+      "FN_4.jpg                  | 474.16         | 419.76         | ✅ YES     \n",
+      "FN_5.jpg                  | 796.23         | 229.17         | ✅ YES     \n",
+      "FN_7.jpg                  | 506.63         | 211.40         | ✅ YES     \n",
+      "FN_8.jpg                  | 432.50         | 237.60         | ✅ YES     \n",
+      "FN_9.jpg                  | 725.21         | 239.89         | ✅ YES     \n",
+      "RB_10.jpg                 | 480.74         | 220.76         | ✅ YES     \n",
+      "RB_11.jpg                 | 449.81         | 224.74         | ✅ YES     \n",
+      "RB_13.jpg                 | 478.89         | 221.27         | ✅ YES     \n",
+      "RB_14.jpg                 | 500.14         | 222.87         | ✅ YES     \n",
+      "RB_2.jpg                  | 494.09         | 223.37         | ✅ YES     \n",
+      "RB_3.jpg                  | 488.29         | 220.04         | ✅ YES     \n",
+      "RB_4.jpg                  | 498.87         | 233.70         | ✅ YES     \n",
+      "RB_6.jpg                  | 496.31         | 221.88         | ✅ YES     \n",
+      "RB_7.jpg                  | 477.51         | 214.55         | ✅ YES     \n",
+      "RB_8.jpg                  | 503.84         | 217.26         | ✅ YES     \n",
+      "RC_10.jpg                 | 491.95         | 221.90         | ✅ YES     \n",
+      "RC_11.jpg                 | 497.75         | 218.97         | ✅ YES     \n",
+      "RC_2.jpg                  | 493.97         | 222.79         | ✅ YES     \n",
+      "RC_3.jpg                  | 481.42         | 220.67         | ✅ YES     \n",
+      "RC_4.jpg                  | 487.19         | 224.59         | ✅ YES     \n",
+      "RC_5.jpg                  | 484.33         | 219.35         | ✅ YES     \n",
+      "RC_6.jpg                  | 483.71         | 216.13         | ✅ YES     \n",
+      "RC_7.jpg                  | 572.23         | 213.54         | ✅ YES     \n",
+      "RC_8.jpg                  | 484.69         | 218.72         | ✅ YES     \n",
+      "RC_9.jpg                  | 578.70         | 218.05         | ✅ YES     \n",
+      "RN_1.jpg                  | 493.88         | 221.69         | ✅ YES     \n",
+      "RN_10.jpg                 | 469.91         | 213.39         | ✅ YES     \n",
+      "RN_11.jpg                 | 488.03         | 212.66         | ✅ YES     \n",
+      "RN_12.jpg                 | 495.62         | 209.92         | ✅ YES     \n",
+      "RN_2.jpg                  | 484.84         | 214.93         | ✅ YES     \n",
+      "RN_3.jpg                  | 467.09         | 218.67         | ✅ YES     \n",
+      "RN_4.jpg                  | 496.70         | 215.96         | ✅ YES     \n",
+      "RN_7.jpg                  | 480.62         | 214.66         | ✅ YES     \n",
+      "RN_8.jpg                  | 502.00         | 213.91         | ✅ YES     \n",
+      "RN_9.jpg                  | 466.22         | 217.81         | ✅ YES     \n",
+      "===========================================================================\n",
+      "                      BENCHMARK SUMMARY\n",
+      "===========================================================================\n",
+      "Total Images Evaluated : 60\n",
+      "Prediction Match Rate  : 100.00% (60/60)\n",
+      "PyTorch Avg Latency    : 503.75 ms / image\n",
+      "ONNX Runtime Avg Latency: 219.87 ms / image\n",
+      "Performance Gain       : 56.35% faster with ONNX Runtime\n",
+      "===========================================================================\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import glob\n",
+    "import time\n",
+    "import tempfile\n",
+    "import warnings\n",
+    "import torch\n",
+    "import torchvision.transforms as transforms\n",
+    "import onnxruntime as ort\n",
+    "import numpy as np\n",
+    "import onnx\n",
+    "from onnxconverter_common import float16\n",
+    "from PIL import Image\n",
+    "from dotenv import load_dotenv\n",
+    "from transformers import ConvNextImageProcessor\n",
+    "from huggingface_hub import HfApi, hf_hub_download\n",
+    "\n",
+    "# Import model loader from base_fusion_model.py\n",
+    "from base_fusion_model import load_fusion_model_from_hf\n",
+    "\n",
+    "# Filter out PyTorch tracing and transformers load warnings for clean logs\n",
+    "warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
+    "\n",
+    "# ==========================================\n",
+    "# 1. PREPROCESSING PIPELINE SETUP\n",
+    "# ==========================================\n",
+    "IMG_SIZE = 260\n",
+    "\n",
+    "val_transforms = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "])\n",
+    "\n",
+    "eff_normalize = transforms.Compose([\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(\n",
+    "        mean=[0.485, 0.456, 0.406],\n",
+    "        std=[0.229, 0.224, 0.225]\n",
+    "    )\n",
+    "])\n",
+    "\n",
+    "\n",
+    "def preprocess_image(pil_image: Image.Image, processor: ConvNextImageProcessor):\n",
+    "    resized_image = val_transforms(pil_image)\n",
+    "    pixel_eff = eff_normalize(resized_image).unsqueeze(0)\n",
+    "    enc = processor(images=resized_image, return_tensors='pt')\n",
+    "    pixel_cnx = enc['pixel_values']\n",
+    "\n",
+    "    return pixel_eff, pixel_cnx\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 2. CONVERT & UPLOAD TO HUGGING FACE\n",
+    "# ==========================================\n",
+    "def convert_and_upload_onnx(\n",
+    "    model: torch.nn.Module,\n",
+    "    processor: ConvNextImageProcessor,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str = \"fusion_model.onnx\",\n",
+    "    hf_token: str = None,\n",
+    "):\n",
+    "    model.eval()\n",
+    "\n",
+    "    dummy_pil = Image.new(\"RGB\", (IMG_SIZE, IMG_SIZE), color=\"white\")\n",
+    "    dummy_eff, dummy_cnx = preprocess_image(dummy_pil, processor)\n",
+    "\n",
+    "    with tempfile.TemporaryDirectory() as temp_dir:\n",
+    "        temp_onnx_path = os.path.join(temp_dir, onnx_filename_in_repo)\n",
+    "\n",
+    "        print(\"Exporting Fusion Model to temporary ONNX graph...\")\n",
+    "        \n",
+    "        with warnings.catch_warnings():\n",
+    "            warnings.simplefilter(\"ignore\")\n",
+    "            torch.onnx.export(\n",
+    "                model,\n",
+    "                (dummy_eff, dummy_cnx),\n",
+    "                temp_onnx_path,\n",
+    "                export_params=True,\n",
+    "                opset_version=18,\n",
+    "                do_constant_folding=True,\n",
+    "                input_names=[\"pixel_values_eff\", \"pixel_values_cnx\"],\n",
+    "                output_names=[\"logits\"],\n",
+    "                dynamic_axes={\n",
+    "                    \"pixel_values_eff\": {0: \"batch_size\"},\n",
+    "                    \"pixel_values_cnx\": {0: \"batch_size\"},\n",
+    "                    \"logits\": {0: \"batch_size\"},\n",
+    "                },\n",
+    "                dynamo=False,\n",
+    "            )\n",
+    "        print(\"✅ Base ONNX export successful.\")\n",
+    "\n",
+    "        # Convert ONNX graph weights to FP16 to keep file size ~142 MB\n",
+    "        print(\"Compressing ONNX model to FP16 precision...\")\n",
+    "        onnx_model = onnx.load(temp_onnx_path)\n",
+    "        onnx_model_fp16 = float16.convert_float_to_float16(onnx_model, keep_io_types=True)\n",
+    "        onnx.save(onnx_model_fp16, temp_onnx_path)\n",
+    "        print(\"✅ ONNX graph compressed to FP16 successfully.\")\n",
+    "\n",
+    "        # Upload directly to Hugging Face Hub using token from .env\n",
+    "        api = HfApi(token=hf_token)\n",
+    "        print(f\"Uploading '{onnx_filename_in_repo}' to Hugging Face repo: '{repo_id}'...\")\n",
+    "        \n",
+    "        api.upload_file(\n",
+    "            path_or_fileobj=temp_onnx_path,\n",
+    "            path_in_repo=onnx_filename_in_repo,\n",
+    "            repo_id=repo_id,\n",
+    "            repo_type=\"model\",\n",
+    "        )\n",
+    "        print(f\"🚀 Successfully uploaded FP16 model to https://huggingface.co/{repo_id}/blob/main/{onnx_filename_in_repo}\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 3. BENCHMARK ON REAL TEST IMAGES\n",
+    "# ==========================================\n",
+    "def benchmark_fusion_with_folder(\n",
+    "    pytorch_model: torch.nn.Module,\n",
+    "    processor: ConvNextImageProcessor,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str,\n",
+    "    image_folder: str = r\"D:\\DamageLensAI\\test_images\",\n",
+    "    hf_token: str = None,\n",
+    "):\n",
+    "    print(f\"Downloading '{onnx_filename_in_repo}' from Hugging Face for evaluation...\")\n",
+    "    onnx_path = hf_hub_download(\n",
+    "        repo_id=repo_id, \n",
+    "        filename=onnx_filename_in_repo,\n",
+    "        token=hf_token\n",
+    "    )\n",
+    "    session = ort.InferenceSession(onnx_path, providers=[\"CPUExecutionProvider\"])\n",
+    "\n",
+    "    supported_exts = (\"*.jpg\", \"*.jpeg\", \"*.png\", \"*.bmp\", \"*.webp\")\n",
+    "    image_paths = []\n",
+    "    for ext in supported_exts:\n",
+    "        image_paths.extend(glob.glob(os.path.join(image_folder, ext)))\n",
+    "\n",
+    "    if not image_paths:\n",
+    "        raise FileNotFoundError(\n",
+    "            f\"❌ No valid images found in folder '{image_folder}'. \"\n",
+    "            f\"Please place your test images inside '{image_folder}'.\"\n",
+    "        )\n",
+    "\n",
+    "    print(f\"📸 Found {len(image_paths)} images in '{image_folder}'. Starting benchmark...\")\n",
+    "\n",
+    "    pytorch_latencies = []\n",
+    "    onnx_latencies = []\n",
+    "    matches = 0\n",
+    "\n",
+    "    # Warmup\n",
+    "    first_img = Image.open(sorted(image_paths)[0]).convert(\"RGB\")\n",
+    "    warmup_eff, warmup_cnx = preprocess_image(first_img, processor)\n",
+    "    warmup_eff_np = warmup_eff.numpy()\n",
+    "    warmup_cnx_np = warmup_cnx.numpy()\n",
+    "\n",
+    "    pytorch_model.eval()\n",
+    "    with torch.no_grad():\n",
+    "        _ = pytorch_model(warmup_eff, warmup_cnx)\n",
+    "    session.run(None, {\"pixel_values_eff\": warmup_eff_np, \"pixel_values_cnx\": warmup_cnx_np})\n",
+    "\n",
+    "    # Benchmark loop\n",
+    "    print(\"\\n\" + \"=\" * 75)\n",
+    "    print(f\"{'IMAGE NAME':<25} | {'PYTORCH (ms)':<14} | {'ONNX (ms)':<14} | {'MATCH?':<10}\")\n",
+    "    print(\"=\" * 75)\n",
+    "\n",
+    "    for img_path in sorted(image_paths):\n",
+    "        file_name = os.path.basename(img_path)\n",
+    "        raw_img = Image.open(img_path).convert(\"RGB\")\n",
+    "        \n",
+    "        pixel_values_eff, pixel_values_cnx = preprocess_image(raw_img, processor)\n",
+    "        np_eff = pixel_values_eff.numpy()\n",
+    "        np_cnx = pixel_values_cnx.numpy()\n",
+    "\n",
+    "        # PyTorch Inference\n",
+    "        start_pt = time.perf_counter()\n",
+    "        with torch.no_grad():\n",
+    "            pt_logits = pytorch_model(pixel_values_eff, pixel_values_cnx)\n",
+    "        pt_time = (time.perf_counter() - start_pt) * 1000\n",
+    "        pytorch_latencies.append(pt_time)\n",
+    "        pt_pred = int(torch.argmax(pt_logits, dim=1)[0])\n",
+    "\n",
+    "        # ONNX Runtime Inference\n",
+    "        start_onnx = time.perf_counter()\n",
+    "        onnx_outputs = session.run(\n",
+    "            None, \n",
+    "            {\"pixel_values_eff\": np_eff, \"pixel_values_cnx\": np_cnx}\n",
+    "        )\n",
+    "        onnx_time = (time.perf_counter() - start_onnx) * 1000\n",
+    "        onnx_latencies.append(onnx_time)\n",
+    "        onnx_pred = int(np.argmax(onnx_outputs[0], axis=1)[0])\n",
+    "\n",
+    "        is_match = pt_pred == onnx_pred\n",
+    "        if is_match:\n",
+    "            matches += 1\n",
+    "        match_str = \"✅ YES\" if is_match else \"❌ NO\"\n",
+    "\n",
+    "        print(f\"{file_name:<25} | {pt_time:<14.2f} | {onnx_time:<14.2f} | {match_str:<10}\")\n",
+    "\n",
+    "    avg_pt = np.mean(pytorch_latencies)\n",
+    "    avg_onnx = np.mean(onnx_latencies)\n",
+    "    speedup = ((avg_pt - avg_onnx) / avg_pt) * 100\n",
+    "    accuracy_match_pct = (matches / len(image_paths)) * 100\n",
+    "\n",
+    "    print(\"=\" * 75)\n",
+    "    print(\"                      BENCHMARK SUMMARY\")\n",
+    "    print(\"=\" * 75)\n",
+    "    print(f\"Total Images Evaluated : {len(image_paths)}\")\n",
+    "    print(f\"Prediction Match Rate  : {accuracy_match_pct:.2f}% ({matches}/{len(image_paths)})\")\n",
+    "    print(f\"PyTorch Avg Latency    : {avg_pt:.2f} ms / image\")\n",
+    "    print(f\"ONNX Runtime Avg Latency: {avg_onnx:.2f} ms / image\")\n",
+    "    print(f\"Performance Gain       : {speedup:.2f}% faster with ONNX Runtime\")\n",
+    "    print(\"=\" * 75 + \"\\n\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# MAIN EXECUTION FLOW\n",
+    "# ==========================================\n",
+    "if __name__ == \"__main__\":\n",
+    "    # Load environment variables from .env file\n",
+    "    load_dotenv()\n",
+    "    HF_TOKEN = os.getenv(\"HF_TOKEN\")\n",
+    "\n",
+    "    if not HF_TOKEN:\n",
+    "        print(\"⚠️ Warning: HF_TOKEN not found in .env file. Pushing/pulling private repos may fail.\")\n",
+    "\n",
+    "    NUM_CLASSES = 6\n",
+    "    HF_REPO_ID = \"junaid17/best_fusion_model_fp16\"\n",
+    "    ONNX_FILENAME = \"fusion_model.onnx\"\n",
+    "    TEST_IMAGE_FOLDER = r\"D:\\DamageLensAI\\test_images\"\n",
+    "    CONVNEXT_MODEL_NAME = \"facebook/convnext-small-224\"\n",
+    "\n",
+    "    # Step 1: Initialize ConvNeXt Processor\n",
+    "    processor = ConvNextImageProcessor.from_pretrained(CONVNEXT_MODEL_NAME)\n",
+    "\n",
+    "    # Step 2: Load PyTorch model from Hugging Face\n",
+    "    py_model = load_fusion_model_from_hf(\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        filename=\"best_fusion_model_fp16.pt\",\n",
+    "        num_classes=NUM_CLASSES,\n",
+    "        device=\"cpu\",\n",
+    "    )\n",
+    "\n",
+    "    # Step 3: Convert to FP16 ONNX and push directly to Hugging Face\n",
+    "    convert_and_upload_onnx(\n",
+    "        model=py_model,\n",
+    "        processor=processor,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )\n",
+    "\n",
+    "    # Step 4: Run benchmarking on test image folder\n",
+    "    benchmark_fusion_with_folder(\n",
+    "        pytorch_model=py_model,\n",
+    "        processor=processor,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        image_folder=TEST_IMAGE_FOLDER,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0e10195c",
+   "metadata": {},
+   "source": [
+    "# Resnet Model Conversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "000b4500",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading checkpoint 'car-damage-classifier.pt' from Hugging Face repo 'junaid17/car-damage-classifier'...\n",
+      "✅ ResNet-18 model loaded successfully.\n",
+      "Exporting ResNet-18 Model to temporary FP32 ONNX graph...\n",
+      "✅ Base ONNX export successful.\n",
+      "Compressing ONNX model to FP16 precision...\n",
+      "✅ ONNX graph compressed to FP16 successfully.\n",
+      "Uploading 'car-damage-classifier.onnx' to Hugging Face repo: 'junaid17/car-damage-classifier'...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Processing Files (1 / 1): 100%|██████████| 22.6MB / 22.6MB, 1.05MB/s  \n",
+      "New Data Upload: 100%|██████████| 22.6MB / 22.6MB, 1.05MB/s  \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🚀 Successfully uploaded FP16 ONNX model to https://huggingface.co/junaid17/car-damage-classifier/blob/main/car-damage-classifier.onnx\n",
+      "Downloading 'car-damage-classifier.onnx' from Hugging Face for evaluation...\n",
+      "📸 Found 60 images in 'D:\\DamageLensAI\\test_images'. Starting benchmark...\n",
+      "\n",
+      "===========================================================================\n",
+      "IMAGE NAME                | PYTORCH (ms)   | ONNX (ms)      | MATCH?    \n",
+      "===========================================================================\n",
+      "FB_10.jpg                 | 46.78          | 14.66          | ✅ YES     \n",
+      "FB_11.jpg                 | 56.13          | 15.03          | ✅ YES     \n",
+      "FB_13.jpg                 | 45.83          | 15.16          | ✅ YES     \n",
+      "FB_15.jpg                 | 44.06          | 14.06          | ✅ YES     \n",
+      "FB_16.jpg                 | 54.86          | 15.04          | ✅ YES     \n",
+      "FB_19.jpg                 | 61.18          | 15.78          | ✅ YES     \n",
+      "FB_20.jpg                 | 47.12          | 13.44          | ✅ YES     \n",
+      "FB_4.jpg                  | 49.21          | 17.78          | ✅ YES     \n",
+      "FB_5.jpg                  | 74.55          | 15.02          | ✅ YES     \n",
+      "FB_9.jpg                  | 48.68          | 14.77          | ✅ YES     \n",
+      "FC_1.jpg                  | 54.30          | 14.43          | ✅ YES     \n",
+      "FC_10.jpg                 | 50.06          | 15.70          | ✅ YES     \n",
+      "FC_11.jpg                 | 48.00          | 13.72          | ✅ YES     \n",
+      "FC_12.jpg                 | 42.24          | 14.62          | ✅ YES     \n",
+      "FC_3.jpg                  | 51.02          | 14.97          | ✅ YES     \n",
+      "FC_4.jpg                  | 47.95          | 15.39          | ✅ YES     \n",
+      "FC_6.jpg                  | 46.55          | 14.33          | ✅ YES     \n",
+      "FC_7.jpg                  | 59.15          | 15.21          | ✅ YES     \n",
+      "FC_8.jpg                  | 48.33          | 16.83          | ✅ YES     \n",
+      "FC_9.jpg                  | 45.76          | 15.09          | ✅ YES     \n",
+      "FN_1.jpg                  | 58.81          | 15.92          | ✅ YES     \n",
+      "FN_10.jpg                 | 46.00          | 15.53          | ✅ YES     \n",
+      "FN_11.jpg                 | 56.70          | 16.44          | ✅ YES     \n",
+      "FN_12.jpg                 | 52.64          | 14.61          | ✅ YES     \n",
+      "FN_3.jpg                  | 52.38          | 14.43          | ✅ YES     \n",
+      "FN_4.jpg                  | 50.76          | 17.84          | ✅ YES     \n",
+      "FN_5.jpg                  | 51.05          | 14.57          | ✅ YES     \n",
+      "FN_7.jpg                  | 49.76          | 14.34          | ✅ YES     \n",
+      "FN_8.jpg                  | 45.30          | 16.05          | ✅ YES     \n",
+      "FN_9.jpg                  | 58.23          | 16.13          | ✅ YES     \n",
+      "RB_10.jpg                 | 58.82          | 16.95          | ✅ YES     \n",
+      "RB_11.jpg                 | 86.68          | 19.35          | ✅ YES     \n",
+      "RB_13.jpg                 | 88.30          | 22.82          | ✅ YES     \n",
+      "RB_14.jpg                 | 76.32          | 18.26          | ✅ YES     \n",
+      "RB_2.jpg                  | 63.31          | 17.44          | ✅ YES     \n",
+      "RB_3.jpg                  | 58.23          | 14.74          | ✅ YES     \n",
+      "RB_4.jpg                  | 55.97          | 15.32          | ✅ YES     \n",
+      "RB_6.jpg                  | 157.29         | 16.34          | ✅ YES     \n",
+      "RB_7.jpg                  | 124.51         | 19.32          | ✅ YES     \n",
+      "RB_8.jpg                  | 63.44          | 15.15          | ✅ YES     \n",
+      "RC_10.jpg                 | 79.38          | 15.67          | ✅ YES     \n",
+      "RC_11.jpg                 | 54.87          | 13.93          | ✅ YES     \n",
+      "RC_2.jpg                  | 37.57          | 13.65          | ✅ YES     \n",
+      "RC_3.jpg                  | 42.66          | 13.28          | ✅ YES     \n",
+      "RC_4.jpg                  | 40.99          | 14.31          | ✅ YES     \n",
+      "RC_5.jpg                  | 55.75          | 17.52          | ✅ YES     \n",
+      "RC_6.jpg                  | 47.62          | 13.57          | ✅ YES     \n",
+      "RC_7.jpg                  | 43.82          | 14.02          | ✅ YES     \n",
+      "RC_8.jpg                  | 47.47          | 14.76          | ✅ YES     \n",
+      "RC_9.jpg                  | 50.02          | 14.33          | ✅ YES     \n",
+      "RN_1.jpg                  | 45.03          | 14.66          | ✅ YES     \n",
+      "RN_10.jpg                 | 44.92          | 14.12          | ✅ YES     \n",
+      "RN_11.jpg                 | 47.08          | 13.45          | ✅ YES     \n",
+      "RN_12.jpg                 | 52.43          | 15.71          | ✅ YES     \n",
+      "RN_2.jpg                  | 49.53          | 16.44          | ✅ YES     \n",
+      "RN_3.jpg                  | 43.17          | 14.90          | ✅ YES     \n",
+      "RN_4.jpg                  | 43.36          | 14.52          | ✅ YES     \n",
+      "RN_7.jpg                  | 50.25          | 14.68          | ✅ YES     \n",
+      "RN_8.jpg                  | 44.29          | 13.59          | ✅ YES     \n",
+      "RN_9.jpg                  | 44.30          | 15.32          | ✅ YES     \n",
+      "===========================================================================\n",
+      "                      BENCHMARK SUMMARY\n",
+      "===========================================================================\n",
+      "Total Images Evaluated : 60\n",
+      "Prediction Match Rate  : 100.00% (60/60)\n",
+      "PyTorch Avg Latency    : 55.68 ms / image\n",
+      "ONNX Runtime Avg Latency: 15.42 ms / image\n",
+      "Performance Gain       : 72.31% faster with ONNX Runtime\n",
+      "===========================================================================\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import glob\n",
+    "import time\n",
+    "import tempfile\n",
+    "import warnings\n",
+    "import torch\n",
+    "import torchvision.transforms as transforms\n",
+    "import onnxruntime as ort\n",
+    "import numpy as np\n",
+    "import onnx\n",
+    "from onnxconverter_common import float16\n",
+    "from PIL import Image\n",
+    "from dotenv import load_dotenv\n",
+    "from huggingface_hub import HfApi, hf_hub_download\n",
+    "\n",
+    "# Import ResNet loader\n",
+    "from base_model_resnet import load_resnet_model_from_hf\n",
+    "\n",
+    "# Suppress PyTorch tracing and deprecation warnings\n",
+    "warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
+    "\n",
+    "# ==========================================\n",
+    "# 1. PREPROCESSING PIPELINE\n",
+    "# ==========================================\n",
+    "test_transforms = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(\n",
+    "        mean=[0.485, 0.456, 0.406],\n",
+    "        std=[0.229, 0.224, 0.225]\n",
+    "    )\n",
+    "])\n",
+    "\n",
+    "\n",
+    "def preprocess_image(pil_image: Image.Image) -> torch.Tensor:\n",
+    "    \"\"\"\n",
+    "    Applies standard ResNet preprocessing and adds a batch dimension.\n",
+    "    Shape returned: (1, 3, 224, 224)\n",
+    "    \"\"\"\n",
+    "    return test_transforms(pil_image).unsqueeze(0)\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 2. CONVERT TO FP16 ONNX & UPLOAD\n",
+    "# ==========================================\n",
+    "def convert_and_upload_onnx_fp16(\n",
+    "    model: torch.nn.Module,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str = \"car-damage-classifier.onnx\",\n",
+    "    hf_token: str = None,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    Exports PyTorch ResNet model to ONNX, compresses it to FP16,\n",
+    "    and pushes it directly to Hugging Face Hub.\n",
+    "    \"\"\"\n",
+    "    model.eval()\n",
+    "\n",
+    "    # Create dummy tensor for graph tracing (1, 3, 224, 224)\n",
+    "    dummy_input = torch.randn(1, 3, 224, 224)\n",
+    "\n",
+    "    with tempfile.TemporaryDirectory() as temp_dir:\n",
+    "        temp_onnx_path = os.path.join(temp_dir, onnx_filename_in_repo)\n",
+    "\n",
+    "        print(\"Exporting ResNet-18 Model to temporary FP32 ONNX graph...\")\n",
+    "        \n",
+    "        with warnings.catch_warnings():\n",
+    "            warnings.simplefilter(\"ignore\")\n",
+    "            torch.onnx.export(\n",
+    "                model,\n",
+    "                dummy_input,\n",
+    "                temp_onnx_path,\n",
+    "                export_params=True,\n",
+    "                opset_version=18,\n",
+    "                do_constant_folding=True,\n",
+    "                input_names=[\"pixel_values\"],\n",
+    "                output_names=[\"logits\"],\n",
+    "                dynamic_axes={\n",
+    "                    \"pixel_values\": {0: \"batch_size\"},\n",
+    "                    \"logits\": {0: \"batch_size\"},\n",
+    "                },\n",
+    "                dynamo=False,\n",
+    "            )\n",
+    "        print(\"✅ Base ONNX export successful.\")\n",
+    "\n",
+    "        # Convert ONNX graph weights to FP16 to reduce file size\n",
+    "        print(\"Compressing ONNX model to FP16 precision...\")\n",
+    "        onnx_model = onnx.load(temp_onnx_path)\n",
+    "        onnx_model_fp16 = float16.convert_float_to_float16(onnx_model, keep_io_types=True)\n",
+    "        onnx.save(onnx_model_fp16, temp_onnx_path)\n",
+    "        print(\"✅ ONNX graph compressed to FP16 successfully.\")\n",
+    "\n",
+    "        # Upload directly to Hugging Face Hub\n",
+    "        api = HfApi(token=hf_token)\n",
+    "        print(f\"Uploading '{onnx_filename_in_repo}' to Hugging Face repo: '{repo_id}'...\")\n",
+    "        \n",
+    "        api.upload_file(\n",
+    "            path_or_fileobj=temp_onnx_path,\n",
+    "            path_in_repo=onnx_filename_in_repo,\n",
+    "            repo_id=repo_id,\n",
+    "            repo_type=\"model\",\n",
+    "        )\n",
+    "        print(f\"🚀 Successfully uploaded FP16 ONNX model to https://huggingface.co/{repo_id}/blob/main/{onnx_filename_in_repo}\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 3. BENCHMARK ON REAL TEST IMAGES\n",
+    "# ==========================================\n",
+    "def benchmark_resnet_with_folder(\n",
+    "    pytorch_model: torch.nn.Module,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str,\n",
+    "    image_folder: str = r\"D:\\DamageLensAI\\test_images\",\n",
+    "    hf_token: str = None,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    Downloads ONNX model from Hugging Face Hub and evaluates PyTorch vs ONNX Runtime performance.\n",
+    "    \"\"\"\n",
+    "    print(f\"Downloading '{onnx_filename_in_repo}' from Hugging Face for evaluation...\")\n",
+    "    onnx_path = hf_hub_download(\n",
+    "        repo_id=repo_id, \n",
+    "        filename=onnx_filename_in_repo,\n",
+    "        token=hf_token\n",
+    "    )\n",
+    "    session = ort.InferenceSession(onnx_path, providers=[\"CPUExecutionProvider\"])\n",
+    "\n",
+    "    supported_exts = (\"*.jpg\", \"*.jpeg\", \"*.png\", \"*.bmp\", \"*.webp\")\n",
+    "    image_paths = []\n",
+    "    for ext in supported_exts:\n",
+    "        image_paths.extend(glob.glob(os.path.join(image_folder, ext)))\n",
+    "\n",
+    "    if not image_paths:\n",
+    "        raise FileNotFoundError(\n",
+    "            f\"❌ No valid images found in folder '{image_folder}'. \"\n",
+    "            f\"Please verify your path.\"\n",
+    "        )\n",
+    "\n",
+    "    print(f\"📸 Found {len(image_paths)} images in '{image_folder}'. Starting benchmark...\")\n",
+    "\n",
+    "    pytorch_latencies = []\n",
+    "    onnx_latencies = []\n",
+    "    matches = 0\n",
+    "\n",
+    "    # Warmup run\n",
+    "    first_img = Image.open(sorted(image_paths)[0]).convert(\"RGB\")\n",
+    "    warmup_tensor = preprocess_image(first_img)\n",
+    "    warmup_np = warmup_tensor.numpy()\n",
+    "\n",
+    "    pytorch_model.eval()\n",
+    "    with torch.no_grad():\n",
+    "        _ = pytorch_model(warmup_tensor)\n",
+    "    session.run(None, {\"pixel_values\": warmup_np})\n",
+    "\n",
+    "    # Benchmark Loop\n",
+    "    print(\"\\n\" + \"=\" * 75)\n",
+    "    print(f\"{'IMAGE NAME':<25} | {'PYTORCH (ms)':<14} | {'ONNX (ms)':<14} | {'MATCH?':<10}\")\n",
+    "    print(\"=\" * 75)\n",
+    "\n",
+    "    for img_path in sorted(image_paths):\n",
+    "        file_name = os.path.basename(img_path)\n",
+    "        raw_img = Image.open(img_path).convert(\"RGB\")\n",
+    "        \n",
+    "        img_tensor = preprocess_image(raw_img)\n",
+    "        np_input = img_tensor.numpy()\n",
+    "\n",
+    "        # --- PyTorch Inference ---\n",
+    "        start_pt = time.perf_counter()\n",
+    "        with torch.no_grad():\n",
+    "            pt_logits = pytorch_model(img_tensor)\n",
+    "        pt_time = (time.perf_counter() - start_pt) * 1000\n",
+    "        pytorch_latencies.append(pt_time)\n",
+    "        pt_pred = int(torch.argmax(pt_logits, dim=1)[0])\n",
+    "\n",
+    "        # --- ONNX Runtime Inference ---\n",
+    "        start_onnx = time.perf_counter()\n",
+    "        onnx_outputs = session.run(None, {\"pixel_values\": np_input})\n",
+    "        onnx_time = (time.perf_counter() - start_onnx) * 1000\n",
+    "        onnx_latencies.append(onnx_time)\n",
+    "        onnx_pred = int(np.argmax(onnx_outputs[0], axis=1)[0])\n",
+    "\n",
+    "        is_match = pt_pred == onnx_pred\n",
+    "        if is_match:\n",
+    "            matches += 1\n",
+    "        match_str = \"✅ YES\" if is_match else \"❌ NO\"\n",
+    "\n",
+    "        print(f\"{file_name:<25} | {pt_time:<14.2f} | {onnx_time:<14.2f} | {match_str:<10}\")\n",
+    "\n",
+    "    avg_pt = np.mean(pytorch_latencies)\n",
+    "    avg_onnx = np.mean(onnx_latencies)\n",
+    "    speedup = ((avg_pt - avg_onnx) / avg_pt) * 100\n",
+    "    accuracy_match_pct = (matches / len(image_paths)) * 100\n",
+    "\n",
+    "    print(\"=\" * 75)\n",
+    "    print(\"                      BENCHMARK SUMMARY\")\n",
+    "    print(\"=\" * 75)\n",
+    "    print(f\"Total Images Evaluated : {len(image_paths)}\")\n",
+    "    print(f\"Prediction Match Rate  : {accuracy_match_pct:.2f}% ({matches}/{len(image_paths)})\")\n",
+    "    print(f\"PyTorch Avg Latency    : {avg_pt:.2f} ms / image\")\n",
+    "    print(f\"ONNX Runtime Avg Latency: {avg_onnx:.2f} ms / image\")\n",
+    "    print(f\"Performance Gain       : {speedup:.2f}% faster with ONNX Runtime\")\n",
+    "    print(\"=\" * 75 + \"\\n\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# MAIN EXECUTION FLOW\n",
+    "# ==========================================\n",
+    "if __name__ == \"__main__\":\n",
+    "    # Load token from .env\n",
+    "    load_dotenv()\n",
+    "    HF_TOKEN = os.getenv(\"HF_TOKEN\")\n",
+    "\n",
+    "    if not HF_TOKEN:\n",
+    "        print(\"⚠️ Warning: HF_TOKEN not found in .env file. Pushing/pulling private repos may fail.\")\n",
+    "\n",
+    "    NUM_CLASSES = 6\n",
+    "    HF_REPO_ID = \"junaid17/car-damage-classifier\"\n",
+    "    MODEL_FILENAME = \"car-damage-classifier.pt\"\n",
+    "    ONNX_FILENAME = \"car-damage-classifier.onnx\"\n",
+    "    TEST_IMAGE_FOLDER = r\"D:\\DamageLensAI\\test_images\"\n",
+    "\n",
+    "    # Step 1: Load PyTorch ResNet Model from Hugging Face\n",
+    "    py_model = load_resnet_model_from_hf(\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        filename=MODEL_FILENAME,\n",
+    "        num_classes=NUM_CLASSES,\n",
+    "        device=\"cpu\",\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )\n",
+    "\n",
+    "    # Step 2: Convert to FP16 ONNX & Push directly to Hugging Face\n",
+    "    convert_and_upload_onnx_fp16(\n",
+    "        model=py_model,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )\n",
+    "\n",
+    "    # Step 3: Run benchmarking on test image folder\n",
+    "    benchmark_resnet_with_folder(\n",
+    "        pytorch_model=py_model,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        image_folder=TEST_IMAGE_FOLDER,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0c387aa2",
+   "metadata": {},
+   "source": [
+    "# YOLO Model Conversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "f0e5e525",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading YOLO checkpoint 'damage_detector.pt' from Hugging Face repo 'junaid17/Yolo_Model'...\n",
+      "✅ YOLO model loaded successfully.\n",
+      "Exporting YOLO11 Model directly to FP16 ONNX format (imgsz=640)...\n",
+      "WARNING 'half' is deprecated and will be removed in the future. Use 'quantize' instead.\n",
+      "Ultralytics 8.4.116  Python-3.11.0 torch-2.13.0+cpu CPU (12th Gen Intel Core i3-1215U)\n",
+      "YOLO11l summary (fused): 190 layers, 25,280,083 parameters, 0 gradients, 86.8 GFLOPs\n",
+      "\n",
+      "\u001b[34m\u001b[1mPyTorch:\u001b[0m starting from 'C:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--Yolo_Model\\snapshots\\0cbc6f112ebc6cb7df05fe93f3f0175cc847bec0\\damage_detector.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 5, 8400) (48.8 MB)\n",
+      "\n",
+      "\u001b[34m\u001b[1mONNX:\u001b[0m starting export with onnx 1.22.0 opset 17...\n",
+      "\u001b[34m\u001b[1mONNX:\u001b[0m slimming with onnxslim 0.1.95...\n",
+      "\u001b[34m\u001b[1mONNX:\u001b[0m converting to FP16...\n",
+      "\u001b[34m\u001b[1mONNX:\u001b[0m export success  7.1s, saved as 'C:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--Yolo_Model\\snapshots\\0cbc6f112ebc6cb7df05fe93f3f0175cc847bec0\\damage_detector.onnx' (48.5 MB)\n",
+      "\n",
+      "Export complete (11.6s)\n",
+      "Results saved to \u001b[1mC:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--Yolo_Model\\snapshots\\0cbc6f112ebc6cb7df05fe93f3f0175cc847bec0\\damage_detector.onnx\u001b[0m\n",
+      "Predict:         yolo predict task=detect model=C:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--Yolo_Model\\snapshots\\0cbc6f112ebc6cb7df05fe93f3f0175cc847bec0\\damage_detector.onnx imgsz=640 quantize=16\n",
+      "Validate:        yolo val task=detect model=C:\\Users\\junai\\.cache\\huggingface\\hub\\models--junaid17--Yolo_Model\\snapshots\\0cbc6f112ebc6cb7df05fe93f3f0175cc847bec0\\damage_detector.onnx imgsz=640 data=/content/drive/MyDrive/custom_data/dataset_custom.yaml quantize=16 \n",
+      "Visualize:       https://netron.app\n",
+      "✅ FP16 ONNX export successful.\n",
+      "📦 Model File Size: 48.50 MB\n",
+      "Uploading 'damage_detector.onnx' to Hugging Face repo: 'junaid17/Yolo_Model'...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Processing Files (1 / 1): 100%|██████████| 50.9MB / 50.9MB, 2.79MB/s  \n",
+      "New Data Upload: 100%|██████████| 8.84MB / 8.84MB,  568kB/s  \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🚀 Successfully uploaded ONNX model to https://huggingface.co/junaid17/Yolo_Model/blob/main/damage_detector.onnx\n",
+      "Downloading 'damage_detector.onnx' from Hugging Face for evaluation...\n",
+      "📸 Found 60 images in 'D:\\DamageLensAI\\test_images'. Starting benchmark...\n",
+      "\n",
+      "===========================================================================\n",
+      "IMAGE NAME                | PYTORCH (ms)   | ONNX (ms)      | MATCH?    \n",
+      "===========================================================================\n",
+      "FB_10.jpg                 | 1961.16        | 933.21         | ✅ YES     \n",
+      "FB_11.jpg                 | 1645.58        | 821.06         | ✅ YES     \n",
+      "FB_13.jpg                 | 1973.07        | 884.90         | ✅ YES     \n",
+      "FB_15.jpg                 | 2296.55        | 1064.76        | ✅ YES     \n",
+      "FB_16.jpg                 | 1159.60        | 679.19         | ✅ YES     \n",
+      "FB_19.jpg                 | 1310.77        | 729.79         | ✅ YES     \n",
+      "FB_20.jpg                 | 1263.96        | 678.22         | ✅ YES     \n",
+      "FB_4.jpg                  | 1385.82        | 658.67         | ✅ YES     \n",
+      "FB_5.jpg                  | 1352.15        | 665.79         | ✅ YES     \n",
+      "FB_9.jpg                  | 1281.14        | 654.91         | ✅ YES     \n",
+      "FC_1.jpg                  | 1430.82        | 639.33         | ✅ YES     \n",
+      "FC_10.jpg                 | 1223.63        | 772.76         | ✅ YES     \n",
+      "FC_11.jpg                 | 1229.15        | 665.59         | ✅ YES     \n",
+      "FC_12.jpg                 | 1442.46        | 654.07         | ✅ YES     \n",
+      "FC_3.jpg                  | 1255.18        | 635.24         | ✅ YES     \n",
+      "FC_4.jpg                  | 1381.86        | 646.25         | ✅ YES     \n",
+      "FC_6.jpg                  | 1380.10        | 644.81         | ✅ YES     \n",
+      "FC_7.jpg                  | 1211.48        | 644.08         | ✅ YES     \n",
+      "FC_8.jpg                  | 1201.65        | 761.60         | ✅ YES     \n",
+      "FC_9.jpg                  | 1429.97        | 644.52         | ✅ YES     \n",
+      "FN_1.jpg                  | 1181.08        | 649.43         | ✅ YES     \n",
+      "FN_10.jpg                 | 1190.17        | 637.96         | ✅ YES     \n",
+      "FN_11.jpg                 | 1590.08        | 764.77         | ✅ YES     \n",
+      "FN_12.jpg                 | 1300.76        | 601.92         | ✅ YES     \n",
+      "FN_3.jpg                  | 1182.89        | 662.93         | ✅ YES     \n",
+      "FN_4.jpg                  | 1638.14        | 650.37         | ✅ YES     \n",
+      "FN_5.jpg                  | 1426.86        | 657.21         | ✅ YES     \n",
+      "FN_7.jpg                  | 1440.60        | 717.18         | ✅ YES     \n",
+      "FN_8.jpg                  | 1385.24        | 621.54         | ✅ YES     \n",
+      "FN_9.jpg                  | 1377.56        | 643.62         | ✅ YES     \n",
+      "RB_10.jpg                 | 1362.92        | 694.52         | ✅ YES     \n",
+      "RB_11.jpg                 | 1352.90        | 630.82         | ✅ YES     \n",
+      "RB_13.jpg                 | 1373.38        | 684.08         | ✅ YES     \n",
+      "RB_14.jpg                 | 1395.17        | 693.95         | ✅ YES     \n",
+      "RB_2.jpg                  | 1323.94        | 619.14         | ✅ YES     \n",
+      "RB_3.jpg                  | 1204.51        | 657.64         | ✅ YES     \n",
+      "RB_4.jpg                  | 1468.20        | 673.22         | ✅ YES     \n",
+      "RB_6.jpg                  | 1317.30        | 635.59         | ✅ YES     \n",
+      "RB_7.jpg                  | 1384.84        | 639.61         | ✅ YES     \n",
+      "RB_8.jpg                  | 1476.26        | 668.84         | ✅ YES     \n",
+      "RC_10.jpg                 | 1292.19        | 650.32         | ✅ YES     \n",
+      "RC_11.jpg                 | 1374.14        | 667.15         | ✅ YES     \n",
+      "RC_2.jpg                  | 1403.93        | 776.57         | ✅ YES     \n",
+      "RC_3.jpg                  | 1314.82        | 644.49         | ✅ YES     \n",
+      "RC_4.jpg                  | 1283.66        | 626.54         | ✅ YES     \n",
+      "RC_5.jpg                  | 1343.69        | 682.90         | ✅ YES     \n",
+      "RC_6.jpg                  | 1348.45        | 603.56         | ✅ YES     \n",
+      "RC_7.jpg                  | 1229.81        | 623.58         | ✅ YES     \n",
+      "RC_8.jpg                  | 1312.53        | 670.69         | ✅ YES     \n",
+      "RC_9.jpg                  | 1196.52        | 713.91         | ✅ YES     \n",
+      "RN_1.jpg                  | 1125.28        | 619.14         | ✅ YES     \n",
+      "RN_10.jpg                 | 1177.20        | 745.49         | ✅ YES     \n",
+      "RN_11.jpg                 | 1377.85        | 632.45         | ✅ YES     \n",
+      "RN_12.jpg                 | 1271.99        | 611.52         | ✅ YES     \n",
+      "RN_2.jpg                  | 1166.40        | 715.83         | ✅ YES     \n",
+      "RN_3.jpg                  | 1209.96        | 602.54         | ✅ YES     \n",
+      "RN_4.jpg                  | 1129.91        | 696.41         | ✅ YES     \n",
+      "RN_7.jpg                  | 1165.09        | 642.48         | ✅ YES     \n",
+      "RN_8.jpg                  | 1327.08        | 600.25         | ✅ YES     \n",
+      "RN_9.jpg                  | 1123.79        | 609.49         | ✅ YES     \n",
+      "===========================================================================\n",
+      "                      BENCHMARK SUMMARY\n",
+      "===========================================================================\n",
+      "Total Images Evaluated : 60\n",
+      "Detection Match Rate   : 100.00% (60/60)\n",
+      "PyTorch Avg Latency    : 1356.05 ms / image\n",
+      "ONNX Runtime Avg Latency: 680.31 ms / image\n",
+      "Performance Gain       : 49.83% faster with ONNX Runtime\n",
+      "===========================================================================\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import glob\n",
+    "import time\n",
+    "import warnings\n",
+    "import torch\n",
+    "import torchvision.transforms as transforms\n",
+    "import onnxruntime as ort\n",
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "from dotenv import load_dotenv\n",
+    "from huggingface_hub import HfApi, hf_hub_download\n",
+    "from ultralytics import YOLO\n",
+    "\n",
+    "# Import YOLO loader\n",
+    "from base_model_yolo import load_yolo_model_from_hf\n",
+    "\n",
+    "# Suppress warnings\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "# ==========================================\n",
+    "# 1. PREPROCESSING PIPELINE\n",
+    "# ==========================================\n",
+    "IMG_SIZE = 640\n",
+    "\n",
+    "yolo_transforms = transforms.Compose([\n",
+    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
+    "    transforms.ToTensor(),\n",
+    "])\n",
+    "\n",
+    "def preprocess_image(pil_image: Image.Image) -> torch.Tensor:\n",
+    "    \"\"\"Preprocesses PIL image for YOLO11 (1, 3, 640, 640).\"\"\"\n",
+    "    return yolo_transforms(pil_image).unsqueeze(0)\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 2. CONVERT, OPTIMIZE TO FP16 & UPLOAD\n",
+    "# ==========================================\n",
+    "def convert_and_upload_yolo_onnx(\n",
+    "    yolo_model: YOLO,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str = \"damage_detector.onnx\",\n",
+    "    hf_token: str = None,\n",
+    "    imgsz: int = 640,\n",
+    "):\n",
+    "    print(f\"Exporting YOLO11 Model directly to FP16 ONNX format (imgsz={imgsz})...\")\n",
+    "    \n",
+    "    # Native Ultralytics FP16 export - lightning fast and avoids converter bugs\n",
+    "    exported_onnx_path = yolo_model.export(\n",
+    "        format=\"onnx\",\n",
+    "        imgsz=imgsz,\n",
+    "        half=True,       # Shrinks the model to ~50MB natively\n",
+    "        dynamic=False,   # Disabling dynamic fixes the CPU type-mismatch errors\n",
+    "        simplify=True,\n",
+    "        opset=17\n",
+    "    )\n",
+    "    print(\"✅ FP16 ONNX export successful.\")\n",
+    "\n",
+    "    file_size_mb = os.path.getsize(exported_onnx_path) / (1024 * 1024)\n",
+    "    print(f\"📦 Model File Size: {file_size_mb:.2f} MB\")\n",
+    "\n",
+    "    # Upload to Hugging Face Hub\n",
+    "    api = HfApi(token=hf_token)\n",
+    "    print(f\"Uploading '{onnx_filename_in_repo}' to Hugging Face repo: '{repo_id}'...\")\n",
+    "    \n",
+    "    api.upload_file(\n",
+    "        path_or_fileobj=exported_onnx_path,\n",
+    "        path_in_repo=onnx_filename_in_repo,\n",
+    "        repo_id=repo_id,\n",
+    "        repo_type=\"model\",\n",
+    "    )\n",
+    "    print(f\"🚀 Successfully uploaded ONNX model to https://huggingface.co/{repo_id}/blob/main/{onnx_filename_in_repo}\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# 3. BENCHMARK ON REAL TEST IMAGES\n",
+    "# ==========================================\n",
+    "def benchmark_yolo_with_folder(\n",
+    "    yolo_model: YOLO,\n",
+    "    repo_id: str,\n",
+    "    onnx_filename_in_repo: str,\n",
+    "    image_folder: str = r\"D:\\DamageLensAI\\test_images\",\n",
+    "    hf_token: str = None,\n",
+    "    conf_thresh: float = 0.25,\n",
+    "):\n",
+    "    print(f\"Downloading '{onnx_filename_in_repo}' from Hugging Face for evaluation...\")\n",
+    "    \n",
+    "    # force_download=True ensures we grab the newly uploaded 50MB file\n",
+    "    onnx_path = hf_hub_download(\n",
+    "        repo_id=repo_id, \n",
+    "        filename=onnx_filename_in_repo,\n",
+    "        token=hf_token,\n",
+    "        force_download=True \n",
+    "    )\n",
+    "    \n",
+    "    session = ort.InferenceSession(onnx_path, providers=[\"CPUExecutionProvider\"])\n",
+    "    input_name = session.get_inputs()[0].name\n",
+    "\n",
+    "    supported_exts = (\"*.jpg\", \"*.jpeg\", \"*.png\", \"*.bmp\", \"*.webp\")\n",
+    "    image_paths = []\n",
+    "    for ext in supported_exts:\n",
+    "        image_paths.extend(glob.glob(os.path.join(image_folder, ext)))\n",
+    "\n",
+    "    if not image_paths:\n",
+    "        raise FileNotFoundError(f\"❌ No valid images found in folder '{image_folder}'.\")\n",
+    "\n",
+    "    print(f\"📸 Found {len(image_paths)} images in '{image_folder}'. Starting benchmark...\")\n",
+    "\n",
+    "    pytorch_latencies = []\n",
+    "    onnx_latencies = []\n",
+    "    matches = 0\n",
+    "\n",
+    "    # Warmup\n",
+    "    first_img = Image.open(sorted(image_paths)[0]).convert(\"RGB\")\n",
+    "    warmup_tensor = preprocess_image(first_img)\n",
+    "    warmup_np = warmup_tensor.numpy()\n",
+    "\n",
+    "    _ = yolo_model.model(warmup_tensor)\n",
+    "    session.run(None, {input_name: warmup_np})\n",
+    "\n",
+    "    print(\"\\n\" + \"=\" * 75)\n",
+    "    print(f\"{'IMAGE NAME':<25} | {'PYTORCH (ms)':<14} | {'ONNX (ms)':<14} | {'MATCH?':<10}\")\n",
+    "    print(\"=\" * 75)\n",
+    "\n",
+    "    for img_path in sorted(image_paths):\n",
+    "        file_name = os.path.basename(img_path)\n",
+    "        raw_img = Image.open(img_path).convert(\"RGB\")\n",
+    "        \n",
+    "        img_tensor = preprocess_image(raw_img)\n",
+    "        np_input = img_tensor.numpy()\n",
+    "\n",
+    "        start_pt = time.perf_counter()\n",
+    "        with torch.no_grad():\n",
+    "            pt_out = yolo_model.model(img_tensor)\n",
+    "        pt_time = (time.perf_counter() - start_pt) * 1000\n",
+    "        pytorch_latencies.append(pt_time)\n",
+    "\n",
+    "        start_onnx = time.perf_counter()\n",
+    "        onnx_out = session.run(None, {input_name: np_input})\n",
+    "        onnx_time = (time.perf_counter() - start_onnx) * 1000\n",
+    "        onnx_latencies.append(onnx_time)\n",
+    "\n",
+    "        pt_preds = pt_out[0] if isinstance(pt_out, (list, tuple)) else pt_out\n",
+    "        pt_max_conf = float(torch.max(pt_preds[0, 4, :])) if pt_preds.ndim == 3 else 0.0\n",
+    "        \n",
+    "        onnx_preds = onnx_out[0]\n",
+    "        onnx_max_conf = float(np.max(onnx_preds[0, 4, :])) if onnx_preds.ndim == 3 else 0.0\n",
+    "\n",
+    "        pt_detected = pt_max_conf >= conf_thresh\n",
+    "        onnx_detected = onnx_max_conf >= conf_thresh\n",
+    "        is_match = (pt_detected == onnx_detected)\n",
+    "\n",
+    "        if is_match:\n",
+    "            matches += 1\n",
+    "        match_str = \"✅ YES\" if is_match else \"❌ NO\"\n",
+    "\n",
+    "        print(f\"{file_name:<25} | {pt_time:<14.2f} | {onnx_time:<14.2f} | {match_str:<10}\")\n",
+    "\n",
+    "    avg_pt = np.mean(pytorch_latencies)\n",
+    "    avg_onnx = np.mean(onnx_latencies)\n",
+    "    speedup = ((avg_pt - avg_onnx) / avg_pt) * 100\n",
+    "    accuracy_match_pct = (matches / len(image_paths)) * 100\n",
+    "\n",
+    "    print(\"=\" * 75)\n",
+    "    print(\"                      BENCHMARK SUMMARY\")\n",
+    "    print(\"=\" * 75)\n",
+    "    print(f\"Total Images Evaluated : {len(image_paths)}\")\n",
+    "    print(f\"Detection Match Rate   : {accuracy_match_pct:.2f}% ({matches}/{len(image_paths)})\")\n",
+    "    print(f\"PyTorch Avg Latency    : {avg_pt:.2f} ms / image\")\n",
+    "    print(f\"ONNX Runtime Avg Latency: {avg_onnx:.2f} ms / image\")\n",
+    "    print(f\"Performance Gain       : {speedup:.2f}% faster with ONNX Runtime\")\n",
+    "    print(\"=\" * 75 + \"\\n\")\n",
+    "\n",
+    "\n",
+    "# ==========================================\n",
+    "# MAIN EXECUTION FLOW\n",
+    "# ==========================================\n",
+    "if __name__ == \"__main__\":\n",
+    "    load_dotenv()\n",
+    "    HF_TOKEN = os.getenv(\"HF_TOKEN\")\n",
+    "\n",
+    "    if not HF_TOKEN:\n",
+    "        print(\"⚠️ Warning: HF_TOKEN not found in .env file.\")\n",
+    "\n",
+    "    HF_REPO_ID = \"junaid17/Yolo_Model\"\n",
+    "    MODEL_FILENAME = \"damage_detector.pt\"\n",
+    "    ONNX_FILENAME = \"damage_detector.onnx\"\n",
+    "    TEST_IMAGE_FOLDER = r\"D:\\DamageLensAI\\test_images\"\n",
+    "\n",
+    "    yolo_model = load_yolo_model_from_hf(\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        filename=MODEL_FILENAME,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )\n",
+    "\n",
+    "    convert_and_upload_yolo_onnx(\n",
+    "        yolo_model=yolo_model,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "        imgsz=IMG_SIZE,\n",
+    "    )\n",
+    "\n",
+    "    benchmark_yolo_with_folder(\n",
+    "        yolo_model=yolo_model,\n",
+    "        repo_id=HF_REPO_ID,\n",
+    "        onnx_filename_in_repo=ONNX_FILENAME,\n",
+    "        image_folder=TEST_IMAGE_FOLDER,\n",
+    "        hf_token=HF_TOKEN,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5b257cbc",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e9c3a35",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8bed307b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "myvenv (3.11.0)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/config.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+import torch
+
+# ---------------- PATHS ----------------
+BASE_DIR = Path(__file__).resolve().parents[1]
+
+DATASET_DIR = BASE_DIR / "data" / "dataset"
+CHECKPOINT_DIR = BASE_DIR / "checkpoints"
+EXPORT_DIR = BASE_DIR / "exports"
+
+CHECKPOINT_DIR.mkdir(exist_ok=True)
+EXPORT_DIR.mkdir(exist_ok=True)
+
+# ---------------- TRAINING ----------------
+BATCH_SIZE = 16
+NUM_WORKERS = 4
+LEARNING_RATE = 1e-4
+WEIGHT_DECAY = 1e-5
+VALIDATION_SPLIT = 0.2
+RANDOM_SEED = 42
+
+# TEMP DEV SETTING
+EPOCHS = 1
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------------- IMAGE SIZES ----------------
+RESNET_IMAGE_SIZE = 224
+FUSION_IMAGE_SIZE = 260
+YOLO_IMAGE_SIZE = 640
+
+# ---------------- YOLO ----------------
+YOLO_BASE_MODEL = "yolo11m.pt"
+YOLO_BATCH_SIZE = 10
+YOLO_EPOCHS = 1
+YOLO_CONFIDENCE_THRESHOLD = 0.05
+
+# ---------------- CLASSES ----------------
+CLASS_NAMES = [
+    "F_Breakage",
+    "F_Crushed",
+    "F_Normal",
+    "R_Breakage",
+    "R_Crushed",
+    "R_Normal"
+]
+
+
+
+CLASS_MAP = {idx: cls for idx, cls in enumerate(CLASS_NAMES)}
+CLASS_TO_IDX = {cls: idx for idx, cls in enumerate(CLASS_NAMES)}
+
+NUM_CLASSES = len(CLASS_NAMES)

src/data/augmentation.py

@@ -0,0 +1,90 @@
+import logging
+from torchvision import transforms
+
+from src.config import RESNET_IMAGE_SIZE, FUSION_IMAGE_SIZE
+
+logger = logging.getLogger(__name__)
+
+
+def get_resnet_train_transforms():
+    logger.info("Creating ResNet training transforms...")
+
+    return transforms.Compose([
+        transforms.Resize((RESNET_IMAGE_SIZE, RESNET_IMAGE_SIZE)),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomRotation(15),
+        transforms.ColorJitter(
+            brightness=0.2,
+            contrast=0.2,
+            saturation=0.2
+        ),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+
+def get_resnet_val_transforms():
+    logger.info("Creating ResNet validation transforms...")
+
+    return transforms.Compose([
+        transforms.Resize((RESNET_IMAGE_SIZE, RESNET_IMAGE_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+
+def get_fusion_train_transforms():
+    logger.info("Creating Fusion training transforms...")
+
+    return transforms.Compose([
+        transforms.Resize((FUSION_IMAGE_SIZE, FUSION_IMAGE_SIZE)),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomRotation(10),
+        transforms.ColorJitter(
+            brightness=0.15,
+            contrast=0.15,
+            saturation=0.15
+        ),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+
+def get_fusion_val_transforms():
+    logger.info("Creating Fusion validation transforms...")
+
+    return transforms.Compose([
+        transforms.Resize((FUSION_IMAGE_SIZE, FUSION_IMAGE_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    resnet_train = get_resnet_train_transforms()
+    resnet_val = get_resnet_val_transforms()
+    fusion_train = get_fusion_train_transforms()
+    fusion_val = get_fusion_val_transforms()
+
+    print("\nTransforms created successfully:")
+    print("ResNet Train:", resnet_train)
+    print("ResNet Val:", resnet_val)
+    print("Fusion Train:", fusion_train)
+    print("Fusion Val:", fusion_val)

src/data/dataset.py

@@ -0,0 +1,189 @@
+import logging
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+from transformers import ConvNextImageProcessor
+
+from src.config import (
+    BATCH_SIZE,
+    NUM_WORKERS
+)
+
+from src.data.ingestion import collect_image_paths
+from src.data.preprocessing import split_dataset
+from src.data.augmentation import (
+    get_resnet_train_transforms,
+    get_resnet_val_transforms,
+    get_fusion_train_transforms,
+    get_fusion_val_transforms
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ResNetDataset(Dataset):
+    def __init__(self, samples, transforms=None):
+        self.samples = samples
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        image_path, label = self.samples[idx]
+
+        image = Image.open(image_path).convert("RGB")
+
+        if self.transforms:
+            image = self.transforms(image)
+
+        return image, label
+
+
+class FusionDataset(Dataset):
+    def __init__(
+        self,
+        samples,
+        transforms=None,
+        convnext_model_name="facebook/convnext-small-224"
+    ):
+        self.samples = samples
+        self.transforms = transforms
+
+        logger.info("Loading ConvNeXt processor...")
+
+        self.processor = ConvNextImageProcessor.from_pretrained(
+            convnext_model_name
+        )
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        image_path, label = self.samples[idx]
+
+        image = Image.open(image_path).convert("RGB")
+
+        if self.transforms:
+            eff_tensor = self.transforms(image)
+        else:
+            raise ValueError("Fusion transforms are required.")
+
+        convnext_inputs = self.processor(
+            images=image,
+            return_tensors="pt"
+        )
+
+        convnext_tensor = convnext_inputs["pixel_values"].squeeze(0)
+
+        return {
+            "pixel_values_eff": eff_tensor,
+            "pixel_values_cnx": convnext_tensor,
+            "labels": label
+        }
+
+
+def create_resnet_dataloaders():
+    logger.info("Creating ResNet dataloaders...")
+
+    samples = collect_image_paths()
+    train_data, val_data = split_dataset(samples)
+
+    train_dataset = ResNetDataset(
+        train_data,
+        transforms=get_resnet_train_transforms()
+    )
+
+    val_dataset = ResNetDataset(
+        val_data,
+        transforms=get_resnet_val_transforms()
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=NUM_WORKERS
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=False,
+        num_workers=NUM_WORKERS
+    )
+
+    logger.info("ResNet dataloaders created successfully.")
+
+    return train_loader, val_loader
+
+
+def create_fusion_dataloaders():
+    logger.info("Creating Fusion dataloaders...")
+
+    samples = collect_image_paths()
+    train_data, val_data = split_dataset(samples)
+
+    train_dataset = FusionDataset(
+        train_data,
+        transforms=get_fusion_train_transforms()
+    )
+
+    val_dataset = FusionDataset(
+        val_data,
+        transforms=get_fusion_val_transforms()
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=NUM_WORKERS
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=False,
+        num_workers=NUM_WORKERS
+    )
+
+    logger.info("Fusion dataloaders created successfully.")
+
+    return train_loader, val_loader
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    print("\nTesting ResNet dataloaders...\n")
+
+    train_loader, val_loader = create_resnet_dataloaders()
+
+    images, labels = next(iter(train_loader))
+
+    print("ResNet batch shape:", images.shape)
+    print("ResNet labels shape:", labels.shape)
+
+    print("\nTesting Fusion dataloaders...\n")
+
+    train_loader, val_loader = create_fusion_dataloaders()
+
+    batch = next(iter(train_loader))
+
+    print(
+        "Fusion EfficientNet batch shape:",
+        batch["pixel_values_eff"].shape
+    )
+
+    print(
+        "Fusion ConvNeXt batch shape:",
+        batch["pixel_values_cnx"].shape
+    )
+
+    print(
+        "Fusion labels shape:",
+        batch["labels"].shape
+    )

src/data/ingestion.py

@@ -0,0 +1,55 @@
+import logging
+from pathlib import Path
+
+from src.config import DATASET_DIR, CLASS_TO_IDX
+
+logger = logging.getLogger(__name__)
+
+VALID_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
+
+
+def collect_image_paths():
+    logger.info("Starting dataset ingestion...")
+
+    if not DATASET_DIR.exists():
+        raise FileNotFoundError(f"Dataset directory not found: {DATASET_DIR}")
+
+    samples = []
+
+    for class_name, label in CLASS_TO_IDX.items():
+        class_dir = DATASET_DIR / class_name
+
+        if not class_dir.exists():
+            logger.warning(f"Missing class folder: {class_dir}")
+            continue
+
+        image_count = 0
+
+        for image_path in class_dir.iterdir():
+            if image_path.suffix.lower() in VALID_EXTENSIONS:
+                samples.append((str(image_path), label))
+                image_count += 1
+
+        logger.info(f"{class_name}: {image_count} images found")
+
+    if not samples:
+        raise ValueError("No valid images found in dataset.")
+
+    logger.info(f"Total images collected: {len(samples)}")
+
+    return samples
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    data = collect_image_paths()
+
+    print(f"\nTotal samples: {len(data)}")
+    print("First 5 samples:")
+
+    for sample in data[:5]:
+        print(sample)

src/data/preprocessing.py

@@ -0,0 +1,58 @@
+import logging
+from collections import Counter
+from sklearn.model_selection import train_test_split
+
+from src.config import VALIDATION_SPLIT, RANDOM_SEED
+from src.data.ingestion import collect_image_paths
+
+logger = logging.getLogger(__name__)
+
+
+def split_dataset(samples):
+    logger.info("Starting dataset preprocessing...")
+
+    if not samples:
+        raise ValueError("Empty dataset provided.")
+
+    image_paths = [sample[0] for sample in samples]
+    labels = [sample[1] for sample in samples]
+
+    logger.info(f"Total samples before split: {len(samples)}")
+
+    train_paths, val_paths, train_labels, val_labels = train_test_split(
+        image_paths,
+        labels,
+        test_size=VALIDATION_SPLIT,
+        stratify=labels,
+        random_state=RANDOM_SEED
+    )
+
+    train_data = list(zip(train_paths, train_labels))
+    val_data = list(zip(val_paths, val_labels))
+
+    logger.info(f"Training samples: {len(train_data)}")
+    logger.info(f"Validation samples: {len(val_data)}")
+
+    logger.info(f"Train distribution: {Counter(train_labels)}")
+    logger.info(f"Validation distribution: {Counter(val_labels)}")
+
+    return train_data, val_data
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    samples = collect_image_paths()
+
+    train_data, val_data = split_dataset(samples)
+
+    print("\nTrain sample preview:")
+    for sample in train_data[:5]:
+        print(sample)
+
+    print("\nValidation sample preview:")
+    for sample in val_data[:5]:
+        print(sample)

src/models/fusion_model.py

@@ -0,0 +1,112 @@
+import logging
+import torch
+import torch.nn as nn
+from torchvision import models
+from transformers import ConvNextModel
+
+logger = logging.getLogger(__name__)
+
+
+class FusionClassifier(nn.Module):
+    def __init__(
+        self,
+        num_classes,
+        convnext_model_name="facebook/convnext-small-224"
+    ):
+        super().__init__()
+
+        logger.info("Initializing Fusion model...")
+
+        # EfficientNet-V2-S
+        eff = models.efficientnet_v2_s(
+            weights=models.EfficientNet_V2_S_Weights.IMAGENET1K_V1
+        )
+
+        for param in eff.parameters():
+            param.requires_grad = False
+
+        for param in eff.features[5].parameters():
+            param.requires_grad = True
+
+        for param in eff.features[6].parameters():
+            param.requires_grad = True
+
+        for param in eff.features[7].parameters():
+            param.requires_grad = True
+
+        self.eff_features = eff.features
+        self.eff_avgpool = eff.avgpool
+        self.eff_out_dim = eff.classifier[1].in_features
+
+        # ConvNeXt
+        cnx = ConvNextModel.from_pretrained(convnext_model_name)
+
+        for param in cnx.parameters():
+            param.requires_grad = False
+
+        for param in cnx.encoder.stages[2].parameters():
+            param.requires_grad = True
+
+        for param in cnx.encoder.stages[3].parameters():
+            param.requires_grad = True
+
+        for param in cnx.layernorm.parameters():
+            param.requires_grad = True
+
+        self.cnx_backbone = cnx
+        self.cnx_out_dim = 768
+
+        fused_dim = self.eff_out_dim + self.cnx_out_dim
+
+        self.fusion_head = nn.Sequential(
+            nn.Dropout(0.4),
+            nn.Linear(fused_dim, 512),
+            nn.LayerNorm(512),
+            nn.GELU(),
+
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+
+            nn.Dropout(0.2),
+            nn.Linear(256, num_classes)
+        )
+
+        logger.info("Fusion model initialized successfully.")
+
+    def forward(self, pixel_values_eff, pixel_values_cnx):
+        x_eff = self.eff_features(pixel_values_eff)
+        x_eff = self.eff_avgpool(x_eff)
+        x_eff = torch.flatten(x_eff, 1)
+
+        cnx_out = self.cnx_backbone(
+            pixel_values=pixel_values_cnx,
+            return_dict=True
+        )
+
+        x_cnx = cnx_out.pooler_output
+
+        fused = torch.cat([x_eff, x_cnx], dim=1)
+
+        logits = self.fusion_head(fused)
+
+        return logits
+
+
+if __name__ == "__main__":
+    import logging
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    model = FusionClassifier(num_classes=6)
+
+    eff_dummy = torch.randn(2, 3, 260, 260)
+    cnx_dummy = torch.randn(2, 3, 224, 224)
+
+    output = model(eff_dummy, cnx_dummy)
+
+    print("Fusion output shape:", output.shape)

src/models/resnet_model.py

@@ -0,0 +1,64 @@
+import logging
+import torch
+import torch.nn as nn
+from torchvision import models
+
+logger = logging.getLogger(__name__)
+
+
+class CarClassifierResNet(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+
+        logger.info("Initializing ResNet18 model...")
+
+        self.model = models.resnet18(weights="DEFAULT")
+
+        # Freeze everything
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        # Unfreeze last layers
+        for param in self.model.layer3.parameters():
+            param.requires_grad = True
+
+        for param in self.model.layer4.parameters():
+            param.requires_grad = True
+
+        # Custom classifier head
+        self.model.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(self.model.fc.in_features, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)
+        )
+
+        logger.info("ResNet18 model initialized successfully.")
+
+    def forward(self, x):
+        return self.model(x)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    model = CarClassifierResNet(num_classes=6)
+
+    dummy_input = torch.randn(2, 3, 128, 128)
+
+    output = model(dummy_input)
+
+    print("Output shape:", output.shape)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(
+        p.numel() for p in model.parameters()
+        if p.requires_grad
+    )
+
+    print("Total params:", total_params)
+    print("Trainable params:", trainable_params)

src/training/train_fusion.py

@@ -0,0 +1,92 @@
+import logging
+import torch.nn as nn
+from torch.optim import AdamW
+
+from src.config import DEVICE, EPOCHS, NUM_CLASSES
+from src.models.fusion_model import FusionClassifier
+from src.data.dataset import create_fusion_dataloaders
+from src.training.trainer import train_dual_input_model
+
+logger = logging.getLogger(__name__)
+
+
+def run_fusion_training():
+    logger.info("Initializing Fusion training pipeline...")
+
+    train_loader, eval_loader = create_fusion_dataloaders()
+
+    model = FusionClassifier(
+        num_classes=NUM_CLASSES
+    ).to(DEVICE)
+
+    criterion = nn.CrossEntropyLoss(
+        label_smoothing=0.1
+    )
+
+    optimizer = AdamW([
+        # EfficientNet unfrozen blocks
+        {
+            "params": model.eff_features[5].parameters(),
+            "lr": 1e-5
+        },
+        {
+            "params": model.eff_features[6].parameters(),
+            "lr": 3e-5
+        },
+        {
+            "params": model.eff_features[7].parameters(),
+            "lr": 3e-5
+        },
+
+        # ConvNeXt unfrozen blocks
+        {
+            "params": model.cnx_backbone.encoder.stages[2].parameters(),
+            "lr": 3e-5
+        },
+        {
+            "params": model.cnx_backbone.encoder.stages[3].parameters(),
+            "lr": 3e-5
+        },
+        {
+            "params": model.cnx_backbone.layernorm.parameters(),
+            "lr": 3e-5
+        },
+
+        # Fusion head
+        {
+            "params": model.fusion_head.parameters(),
+            "lr": 1e-4
+        }
+
+    ], weight_decay=1e-4)
+
+    logger.info("Starting Fusion training...")
+
+    all_preds, all_labels = train_dual_input_model(
+        model=model,
+        train_loader=train_loader,
+        eval_loader=eval_loader,
+        optimizer=optimizer,
+        criterion=criterion,
+        device=DEVICE,
+        epochs=EPOCHS,
+        checkpoint_model_name="best_fusion_model",
+        patience=7
+    )
+
+    logger.info("Fusion training completed.")
+
+    return all_preds, all_labels
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    preds, labels = run_fusion_training()
+
+    print("\nFusion training completed successfully.")
+    print("Prediction samples:", preds[:10])
+    print("Label samples:", labels[:10])

src/training/train_resnet.py

@@ -0,0 +1,68 @@
+import logging
+import torch.nn as nn
+from torch.optim import AdamW
+
+from src.config import DEVICE, EPOCHS, NUM_CLASSES
+from src.models.resnet_model import CarClassifierResNet
+from src.data.dataset import create_resnet_dataloaders
+from src.training.trainer import train_single_input_model
+
+logger = logging.getLogger(__name__)
+
+
+def run_resnet_training():
+    logger.info("Initializing ResNet training pipeline...")
+
+    train_loader, eval_loader = create_resnet_dataloaders()
+
+    model = CarClassifierResNet(
+        num_classes=NUM_CLASSES
+    ).to(DEVICE)
+
+    criterion = nn.CrossEntropyLoss()
+
+    optimizer = AdamW([
+        {
+            "params": model.model.layer3.parameters(),
+            "lr": 1e-5
+        },
+        {
+            "params": model.model.layer4.parameters(),
+            "lr": 1e-5
+        },
+        {
+            "params": model.model.fc.parameters(),
+            "lr": 1e-4
+        }
+    ])
+
+    logger.info("Starting ResNet training...")
+
+    all_preds, all_labels = train_single_input_model(
+        model=model,
+        train_loader=train_loader,
+        eval_loader=eval_loader,
+        optimizer=optimizer,
+        criterion=criterion,
+        device=DEVICE,
+        epochs=EPOCHS,
+        checkpoint_model_name="best_resnet_model",
+        patience=7
+    )
+
+    logger.info("ResNet training completed.")
+
+    return all_preds, all_labels
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    preds, labels = run_resnet_training()
+
+    print("\nTraining completed successfully.")
+    print("Prediction samples:", preds[:10])
+    print("Label samples:", labels[:10])

src/training/train_yolo.py

@@ -0,0 +1,85 @@
+import logging
+from shutil import copy2, rmtree
+from ultralytics import YOLO
+
+from src.config import BASE_DIR, CHECKPOINT_DIR, DEVICE
+
+logger = logging.getLogger(__name__)
+
+YOLO_DATASET_CONFIG = BASE_DIR / "data" / "yolo" / "dataset_custom.yaml"
+YOLO_BASE_MODEL = CHECKPOINT_DIR / "yolo11m.pt"
+
+
+def run_yolo_training():
+    logger.info("Initializing YOLO training pipeline...")
+
+    if not YOLO_DATASET_CONFIG.exists():
+        raise FileNotFoundError(
+            f"YOLO dataset config not found: {YOLO_DATASET_CONFIG}"
+        )
+
+    if not YOLO_BASE_MODEL.exists():
+        raise FileNotFoundError(
+            f"YOLO base model not found: {YOLO_BASE_MODEL}"
+        )
+
+    yolo_device = 0 if DEVICE == "cuda" else "cpu"
+
+    checkpoint_root = CHECKPOINT_DIR.resolve()
+    temp_run_name = "temp_yolo_run"
+
+    logger.info("Loading YOLO base model...")
+
+    model = YOLO(str(YOLO_BASE_MODEL.resolve()))
+
+    logger.info("Starting YOLO training...")
+
+    model.train(
+        data=str(YOLO_DATASET_CONFIG.resolve()),
+        imgsz=416,
+        batch=4,
+        epochs=1,
+        device=yolo_device,
+        project=str(checkpoint_root),
+        name=temp_run_name,
+        exist_ok=True
+    )
+
+    best_model_path = (
+        checkpoint_root /
+        temp_run_name /
+        "weights" /
+        "best.pt"
+    )
+
+    if not best_model_path.exists():
+        raise FileNotFoundError(
+            f"YOLO best model not found: {best_model_path}"
+        )
+
+    final_model_path = checkpoint_root / "damage_detector.pt"
+
+    copy2(best_model_path, final_model_path)
+
+    logger.info(f"Final YOLO model saved at: {final_model_path}")
+
+    # cleanup temp training folder
+    temp_run_dir = checkpoint_root / temp_run_name
+
+    if temp_run_dir.exists():
+        rmtree(temp_run_dir)
+        logger.info("Temporary YOLO training artifacts deleted.")
+
+    return final_model_path
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    model_path = run_yolo_training()
+
+    print("\nYOLO training completed successfully.")
+    print(f"Saved model: {model_path}")

src/training/trainer.py

@@ -0,0 +1,305 @@
+import logging
+import torch
+from tqdm import tqdm
+from transformers import get_cosine_schedule_with_warmup
+
+from src.config import CHECKPOINT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+class EarlyStopping:
+    def __init__(self, patience=7, min_delta=0.001):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+
+    def __call__(self, val_acc):
+        if self.best_score is None:
+            self.best_score = val_acc
+
+        elif val_acc < self.best_score + self.min_delta:
+            self.counter += 1
+
+            logger.info(
+                f"EarlyStopping counter: {self.counter}/{self.patience}"
+            )
+
+            if self.counter >= self.patience:
+                self.early_stop = True
+
+        else:
+            self.best_score = val_acc
+            self.counter = 0
+
+
+def train_single_input_model(
+    model,
+    train_loader,
+    eval_loader,
+    optimizer,
+    criterion,
+    device,
+    epochs,
+    checkpoint_model_name,
+    patience=7
+):
+    logger.info("Starting single-input training...")
+
+    num_training_steps = epochs * len(train_loader)
+    num_warmup_steps = int(0.1 * num_training_steps)
+
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps
+    )
+
+    early_stopping = EarlyStopping(patience=patience)
+
+    best_acc = 0.0
+    all_preds = []
+    all_labels = []
+
+    for epoch in range(epochs):
+        logger.info(f"Epoch {epoch + 1}/{epochs}")
+
+        model.train()
+
+        running_loss = 0
+        correct = 0
+        total = 0
+
+        for images, labels in tqdm(
+            train_loader,
+            desc=f"Epoch {epoch+1} Training"
+        ):
+            images = images.to(device)
+            labels = labels.to(device)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            logits = model(images)
+
+            loss = criterion(logits, labels)
+
+            loss.backward()
+
+            optimizer.step()
+            scheduler.step()
+
+            running_loss += loss.item()
+
+            preds = torch.argmax(logits, dim=1)
+
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+
+        train_loss = running_loss / len(train_loader)
+        train_acc = 100 * correct / total
+
+        model.eval()
+
+        val_running_loss = 0
+        val_correct = 0
+        val_total = 0
+
+        all_preds = []
+        all_labels = []
+
+        with torch.no_grad():
+            for images, labels in tqdm(
+                eval_loader,
+                desc=f"Epoch {epoch+1} Validation"
+            ):
+                images = images.to(device)
+                labels = labels.to(device)
+
+                logits = model(images)
+
+                loss = criterion(logits, labels)
+
+                val_running_loss += loss.item()
+
+                preds = torch.argmax(logits, dim=1)
+
+                val_correct += (preds == labels).sum().item()
+                val_total += labels.size(0)
+
+                all_preds.extend(preds.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        val_loss = val_running_loss / len(eval_loader)
+        val_acc = 100 * val_correct / val_total
+
+        logger.info(
+            f"Train Loss: {train_loss:.4f} | "
+            f"Train Acc: {train_acc:.2f}% || "
+            f"Val Loss: {val_loss:.4f} | "
+            f"Val Acc: {val_acc:.2f}%"
+        )
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+
+            checkpoint_path = CHECKPOINT_DIR / f"{checkpoint_model_name}.pt"
+
+            torch.save(
+                {
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "epoch": epoch,
+                    "val_acc": val_acc
+                },
+                checkpoint_path
+            )
+
+            logger.info(f"Best checkpoint saved at: {checkpoint_path}")
+
+        early_stopping(val_acc)
+
+        if early_stopping.early_stop:
+            logger.info("Early stopping triggered.")
+            break
+
+    return all_preds, all_labels
+
+
+def train_dual_input_model(
+    model,
+    train_loader,
+    eval_loader,
+    optimizer,
+    criterion,
+    device,
+    epochs,
+    checkpoint_model_name,
+    patience=7
+):
+    logger.info("Starting dual-input training...")
+
+    num_training_steps = epochs * len(train_loader)
+    num_warmup_steps = int(0.1 * num_training_steps)
+
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps
+    )
+
+    early_stopping = EarlyStopping(patience=patience)
+
+    best_acc = 0.0
+    all_preds = []
+    all_labels = []
+
+    for epoch in range(epochs):
+        logger.info(f"Epoch {epoch + 1}/{epochs}")
+
+        model.train()
+
+        running_loss = 0
+        correct = 0
+        total = 0
+
+        for batch in tqdm(
+            train_loader,
+            desc=f"Epoch {epoch+1} Training"
+        ):
+            images_eff = batch["pixel_values_eff"].to(device)
+            images_cnx = batch["pixel_values_cnx"].to(device)
+            labels = batch["labels"].to(device)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            logits = model(images_eff, images_cnx)
+
+            loss = criterion(logits, labels)
+
+            loss.backward()
+
+            optimizer.step()
+            scheduler.step()
+
+            running_loss += loss.item()
+
+            preds = torch.argmax(logits, dim=1)
+
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+
+        train_loss = running_loss / len(train_loader)
+        train_acc = 100 * correct / total
+
+        model.eval()
+
+        val_running_loss = 0
+        val_correct = 0
+        val_total = 0
+
+        all_preds = []
+        all_labels = []
+
+        with torch.no_grad():
+            for batch in tqdm(
+                eval_loader,
+                desc=f"Epoch {epoch+1} Validation"
+            ):
+                images_eff = batch["pixel_values_eff"].to(device)
+                images_cnx = batch["pixel_values_cnx"].to(device)
+                labels = batch["labels"].to(device)
+
+                logits = model(images_eff, images_cnx)
+
+                loss = criterion(logits, labels)
+
+                val_running_loss += loss.item()
+
+                preds = torch.argmax(logits, dim=1)
+
+                val_correct += (preds == labels).sum().item()
+                val_total += labels.size(0)
+
+                all_preds.extend(preds.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        val_loss = val_running_loss / len(eval_loader)
+        val_acc = 100 * val_correct / val_total
+
+        logger.info(
+            f"Train Loss: {train_loss:.4f} | "
+            f"Train Acc: {train_acc:.2f}% || "
+            f"Val Loss: {val_loss:.4f} | "
+            f"Val Acc: {val_acc:.2f}%"
+        )
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+
+            checkpoint_path = CHECKPOINT_DIR / f"{checkpoint_model_name}.pt"
+
+            torch.save(
+                {
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "epoch": epoch,
+                    "val_acc": val_acc
+                },
+                checkpoint_path
+            )
+
+            logger.info(f"Best checkpoint saved at: {checkpoint_path}")
+
+        early_stopping(val_acc)
+
+        if early_stopping.early_stop:
+            logger.info("Early stopping triggered.")
+            break
+
+    return all_preds, all_labels
+
+
+if __name__ == "__main__":
+    print("Trainer utilities ready.")

test/test_augmentation.py

@@ -0,0 +1,40 @@
+import logging
+from PIL import Image
+
+from src.data.augmentation import (
+    get_resnet_train_transforms,
+    get_fusion_train_transforms
+)
+
+logger = logging.getLogger(__name__)
+
+
+def test_augmentation():
+    logger.info("Testing augmentation pipelines...")
+
+    dummy_image = Image.new("RGB", (300, 300))
+
+    resnet_transform = get_resnet_train_transforms()
+    fusion_transform = get_fusion_train_transforms()
+
+    resnet_tensor = resnet_transform(dummy_image)
+    fusion_tensor = fusion_transform(dummy_image)
+
+    assert resnet_tensor.shape == (3, 128, 128), \
+        f"Unexpected ResNet shape: {resnet_tensor.shape}"
+
+    assert fusion_tensor.shape == (3, 260, 260), \
+        f"Unexpected Fusion shape: {fusion_tensor.shape}"
+
+    logger.info("Augmentation test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_augmentation()
+
+    print("Augmentation test completed successfully.")

test/test_config.py

@@ -0,0 +1,37 @@
+import logging
+
+from src.config import (
+    BASE_DIR,
+    CHECKPOINT_DIR,
+    DEVICE,
+    BATCH_SIZE,
+    EPOCHS,
+    NUM_CLASSES
+)
+
+logger = logging.getLogger(__name__)
+
+
+def test_config():
+    logger.info("Testing config settings...")
+
+    assert BASE_DIR.exists(), "BASE_DIR missing"
+    assert CHECKPOINT_DIR.exists(), "CHECKPOINT_DIR missing"
+
+    assert DEVICE in ["cpu", "cuda"], "Invalid device"
+    assert BATCH_SIZE > 0, "Invalid batch size"
+    assert EPOCHS > 0, "Invalid epochs"
+    assert NUM_CLASSES == 6, "NUM_CLASSES mismatch"
+
+    logger.info("Config test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_config()
+
+    print("Config test completed successfully.")

test/test_dataset.py

@@ -0,0 +1,57 @@
+import logging
+
+from src.data.dataset import (
+    create_resnet_dataloaders,
+    create_fusion_dataloaders
+)
+
+logger = logging.getLogger(__name__)
+
+
+def test_dataset():
+    logger.info("Testing dataset loaders...")
+
+    # ---------------- ResNet ----------------
+    resnet_loader, _ = create_resnet_dataloaders()
+
+    images, labels = next(iter(resnet_loader))
+
+    assert images.shape[1:] == (3, 128, 128), \
+        f"Unexpected ResNet image shape: {images.shape}"
+
+    assert len(labels.shape) == 1, \
+        f"Unexpected ResNet labels shape: {labels.shape}"
+
+    logger.info("ResNet dataloader test passed.")
+
+    # ---------------- Fusion ----------------
+    fusion_loader, _ = create_fusion_dataloaders()
+
+    batch = next(iter(fusion_loader))
+
+    assert "pixel_values_eff" in batch, "Missing EfficientNet input"
+    assert "pixel_values_cnx" in batch, "Missing ConvNeXt input"
+    assert "labels" in batch, "Missing labels"
+
+    assert batch["pixel_values_eff"].shape[1:] == (3, 260, 260), \
+        f"Unexpected Fusion EfficientNet shape: {batch['pixel_values_eff'].shape}"
+
+    assert batch["pixel_values_cnx"].shape[1:] == (3, 224, 224), \
+        f"Unexpected Fusion ConvNeXt shape: {batch['pixel_values_cnx'].shape}"
+
+    assert len(batch["labels"].shape) == 1, \
+        f"Unexpected Fusion labels shape: {batch['labels'].shape}"
+
+    logger.info("Fusion dataloader test passed.")
+    logger.info("Dataset test passed successfully.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_dataset()
+
+    print("Dataset test completed successfully.")

test/test_fusion_model.py

@@ -0,0 +1,42 @@
+import logging
+import torch
+
+from src.models.fusion_model import FusionClassifier
+from src.config import NUM_CLASSES
+
+logger = logging.getLogger(__name__)
+
+
+def test_fusion_model():
+    logger.info("Testing Fusion model architecture...")
+
+    model = FusionClassifier(
+        num_classes=NUM_CLASSES
+    )
+
+    model.eval()
+
+    eff_dummy = torch.randn(2, 3, 260, 260)
+    cnx_dummy = torch.randn(2, 3, 224, 224)
+
+    with torch.no_grad():
+        output = model(
+            eff_dummy,
+            cnx_dummy
+        )
+
+    assert output.shape == (2, NUM_CLASSES), \
+        f"Unexpected output shape: {output.shape}"
+
+    logger.info("Fusion model test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_fusion_model()
+
+    print("Fusion model test completed successfully.")

test/test_ingestion.py

@@ -0,0 +1,32 @@
+import logging
+import os
+
+from src.data.ingestion import collect_image_paths
+
+logger = logging.getLogger(__name__)
+
+
+def test_ingestion():
+    logger.info("Testing ingestion...")
+
+    samples = collect_image_paths()
+
+    assert len(samples) > 0, "No samples found"
+
+    image_path, label = samples[0]
+
+    assert os.path.exists(image_path), "Image path missing"
+    assert isinstance(label, int), "Label invalid"
+
+    logger.info("Ingestion test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_ingestion()
+
+    print("Ingestion test completed successfully.")

test/test_preprocessing.py

@@ -0,0 +1,37 @@
+import logging
+
+from src.data.ingestion import collect_image_paths
+from src.data.preprocessing import split_dataset
+
+logger = logging.getLogger(__name__)
+
+
+def test_preprocessing():
+    logger.info("Testing preprocessing...")
+
+    samples = collect_image_paths()
+
+    train_data, val_data = split_dataset(samples)
+
+    assert len(train_data) > 0, "Train split is empty"
+    assert len(val_data) > 0, "Validation split is empty"
+
+    train_paths = set(x[0] for x in train_data)
+    val_paths = set(x[0] for x in val_data)
+
+    overlap = train_paths.intersection(val_paths)
+
+    assert len(overlap) == 0, "Train and validation overlap found"
+
+    logger.info("Preprocessing test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_preprocessing()
+
+    print("Preprocessing test completed successfully.")

test/test_resnet_model.py

@@ -0,0 +1,38 @@
+import logging
+import torch
+
+from src.models.resnet_model import CarClassifierResNet
+from src.config import NUM_CLASSES
+
+logger = logging.getLogger(__name__)
+
+
+def test_resnet_model():
+    logger.info("Testing ResNet model architecture...")
+
+    model = CarClassifierResNet(
+        num_classes=NUM_CLASSES
+    )
+
+    model.eval()
+
+    dummy_input = torch.randn(2, 3, 128, 128)
+
+    with torch.no_grad():
+        output = model(dummy_input)
+
+    assert output.shape == (2, NUM_CLASSES), \
+        f"Unexpected output shape: {output.shape}"
+
+    logger.info("ResNet model test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_resnet_model()
+
+    print("ResNet model test completed successfully.")

test/test_train_fusion.py

@@ -0,0 +1,39 @@
+import logging
+
+from src.training.train_fusion import run_fusion_training
+from src.config import CHECKPOINT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def test_train_fusion():
+    logger.info("Testing Fusion training pipeline...")
+
+    checkpoint_path = CHECKPOINT_DIR / "best_fusion_model.pt"
+
+    if checkpoint_path.exists():
+        checkpoint_path.unlink()
+
+    preds, labels = run_fusion_training()
+
+    assert checkpoint_path.exists(), \
+        "Fusion checkpoint was not created"
+
+    assert len(preds) > 0, \
+        "No predictions returned"
+
+    assert len(labels) > 0, \
+        "No labels returned"
+
+    logger.info("Fusion training test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_train_fusion()
+
+    print("Fusion training test completed successfully.")

test/test_train_resnet.py

@@ -0,0 +1,39 @@
+import logging
+
+from src.training.train_resnet import run_resnet_training
+from src.config import CHECKPOINT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def test_train_resnet():
+    logger.info("Testing ResNet training pipeline...")
+
+    checkpoint_path = CHECKPOINT_DIR / "best_resnet_model.pt"
+
+    if checkpoint_path.exists():
+        checkpoint_path.unlink()
+
+    preds, labels = run_resnet_training()
+
+    assert checkpoint_path.exists(), \
+        "ResNet checkpoint was not created"
+
+    assert len(preds) > 0, \
+        "No predictions returned"
+
+    assert len(labels) > 0, \
+        "No labels returned"
+
+    logger.info("ResNet training test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_train_resnet()
+
+    print("ResNet training test completed successfully.")

test/test_train_yolo.py

@@ -0,0 +1,36 @@
+import logging
+
+from src.training.train_yolo import run_yolo_training
+from src.config import CHECKPOINT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def test_train_yolo():
+    logger.info("Testing YOLO training pipeline...")
+
+    checkpoint_path = CHECKPOINT_DIR / "damage_detector.pt"
+
+    if checkpoint_path.exists():
+        checkpoint_path.unlink()
+
+    output_path = run_yolo_training()
+
+    assert checkpoint_path.exists(), \
+        "YOLO checkpoint was not created"
+
+    assert output_path.exists(), \
+        "Returned YOLO model path invalid"
+
+    logger.info("YOLO training test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_train_yolo()
+
+    print("YOLO training test completed successfully.")