Upload 43 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/yolo_detection_sample.jpg filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.10-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+
+WORKDIR /app
+
+# --- SYSTEM DEPENDENCIES (CRITICAL FOR OPENCV / YOLO) ---
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    gcc \
+    libgl1 \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# --- PYTHON DEPENDENCIES ---
+COPY requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip \
+    && pip install --no-cache-dir -r requirements.txt
+
+# --- APP CODE ---
+COPY . .
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,10 +1,702 @@
+# 🚗 DamageLens: AI-Powered Car Damage Detection
+
+[![Python 3.11+](https://img.shields.io/badge/Python-3.11%2B-brightgreen)](https://python.org)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.0%2B-red)](https://pytorch.org)
+[![FastAPI](https://img.shields.io/badge/FastAPI-Latest-teal)](https://fastapi.tiangolo.com)
+[![CI Pipeline](https://github.com/junaidariie/DamageLensAI/actions/workflows/ci.yaml/badge.svg)](https://github.com/junaidariie/DamageLensAI/actions/workflows/ci.yaml)
+[![License](https://img.shields.io/badge/License-MIT-orange)](LICENSE)
+
+---
+
+## ⚠️ Important Notes
+
+> **Cold Startup Time**: The API may take **4-5 minutes** on the first request to warm up the models. Subsequent predictions will be significantly faster.
+
+> **Model Size**: The Fusion model is computationally intensive. Individual predictions typically complete in 30-60 seconds depending on hardware.
+
+---
+
+**APP LINK** : https://junaidariie.github.io/DamageLensAI/
+
+**HF REPO** : https://huggingface.co/spaces/junaid17/DamageLensAI/tree/main
+
+---
+
+## 📋 Table of Contents
+
+- [Overview](#-overview)
+- [Features](#-features)
+- [Architecture](#-architecture)
+- [Model Performance](#-model-performance)
+- [CI Pipeline](#-ci-pipeline)
+- [Setup & Installation](#-setup--installation)
+- [Usage](#-usage)
+- [API Documentation](#-api-documentation)
+- [Model Optimization](#-model-optimization)
+- [Dataset & Training](#-dataset--training)
+- [Web UI Features](#-web-ui-features)
+- [Directory Structure](#-directory-structure)
+- [Limitations & Known Issues](#-limitations--known-issues)
+
+---
+
+## 🎯 Overview
+
+**DamageLens** is an advanced AI system for detecting and classifying car damage using multi-model fusion architecture. It combines the power of **ResNet-18**, **EfficientNet-V2-S**, and **ConvNeXt-Small** to achieve robust damage classification across vehicle front and rear sections.
+
+The system can identify six damage categories:
+- ✅ Front Normal / Front Breakage / Front Crushed
+- ✅ Rear Normal / Rear Breakage / Rear Crushed
+
+Additionally, it uses **YOLO object detection** to localize damage regions with bounding boxes.
+
+---
+
+## ✨ Features
+
+| Feature | Description |
+|---------|-------------|
+| **Dual Model Architecture** | ResNet (lightweight) and Fusion (high-accuracy) options |
+| **Grad-CAM Visualization** | Understand which image regions drive predictions |
+| **Real-time YOLO Detection** | Localize damage with confidence scores |
+| **FP16 Optimization** | Reduced model size (788MB → 135MB) with minimal accuracy loss |
+| **FastAPI Backend** | High-performance REST API with async support |
+| **Responsive Web UI** | Modern, interactive web interface with real-time feedback |
+| **Static File Serving** | Efficient caching and delivery of results |
+| **CI/CD Pipeline** | Automated testing via GitHub Actions on every push/PR |
+| **HuggingFace Integration** | Models auto-downloaded from HF Hub on first startup |
+
+---
+
+## 🏗️ Architecture
+
+### System Overview
+
+```
+┌──────────────────────────────────────────────────────┐
+│                   Frontend (Web UI)                  │
+│  HTML / CSS / JavaScript  (Dark Mode, Glassmorphism) │
+│  ├─ Drag & Drop Image Upload                         │
+│  ├─ Model Selection (Fusion / ResNet)                │
+│  └─ Real-time Result Tabs (Prediction/GradCAM/YOLO)  │
+└───────────────────┬──────────────────────────────────┘
+                    │ REST API (JSON)
+┌───────────────────▼──────────────────────────────────┐
+│              FastAPI Backend  (app.py)               │
+│  ├─ POST /predict/resnet    → ResNet inference       │
+│  ├─ POST /predict/fusion    → Fusion inference       │
+│  ├─ POST /predict?mode=*    → Grad-CAM generation    │
+│  └─ POST /predict/yolo      → YOLO detection         │
+│                                                      │
+│  Lifespan: models loaded once at startup             │
+│  Static:   /static/uploads  /static/results          │
+└──────┬───────────┬──────────────┬────────────────────┘
+       │           │              │
+┌──────▼──┐  ┌─────▼──────┐  ┌───▼──────────┐
+│ ResNet  │  │   Fusion   │  │  YOLO v11m   │
+│  (77%)  │  │   (84%)    │  │  Detection   │
+└──────┬──┘  └─────┬──────┘  └───┬──────────┘
+       │           │              │
+       └─────┬─────┘              │
+             │                    │
+     ┌───────▼──────┐    ┌────────▼────────┐
+     │  Grad-CAM    │    │  Bounding Boxes │
+     │  Heatmaps    │    │  + Confidence   │
+     └──────────────┘    └─────────────────┘
+```
+
+### Model Loading (scripts/load_models.py)
+
+```
+Startup
+  │
+  ├─ hf_hub_download("junaid17/car-damage-classifier")
+  │       └─> ResnetCarDamagePredictor(checkpoint, class_map)
+  │
+  ├─ hf_hub_download("junaid17/best_fusion_model_fp16")
+  │       └─> FusionCarDamagePredictor(checkpoint, class_map)
+  │
+  └─ hf_hub_download("junaid17/Yolo_Model")
+          └─> YOLO(checkpoint)
+```
+
+### Fusion Model (High Accuracy — 84%)
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                          INPUT IMAGE                            │
+│                         (3, 260, 260)                           │
+└────────────────┬────────────────────────────────┬──────────────┘
+                 │                                │
+         ┌───────▼────────┐             ┌─────────▼────────┐
+         │ EfficientNet-  │             │  ConvNeXt-Small  │
+         │ V2-S Backbone  │             │  Backbone        │
+         │                │             │                  │
+         │ Frozen except  │             │ Frozen except    │
+         │ features[5,6,7]│             │ stages[2,3] +    │
+         │ (unfrozen)     │             │ layernorm        │
+         └───────┬────────┘             └─────────┬────────┘
+                 │                                │
+         ┌───────▼────────┐             ┌─────────▼────────┐
+         │ AdaptiveAvg    │             │  Pooler Output   │
+         │ Pool → Flatten │             │                  │
+         └───────┬────────┘             └─────────┬────────┘
+                 │  (1280,)                        │  (768,)
+                 └──────────────┬─────────────────┘
+                                │
+                        ┌───────▼────────┐
+                        │  CONCATENATE   │
+                        │  1280 + 768    │
+                        │  = (2048,)     │
+                        └───────┬────────┘
+                                │
+                    ┌───────────▼───────────┐
+                    │   FUSION HEAD         │
+                    │  Dropout(0.4)         │
+                    │  Linear(2048 → 512)   │
+                    │  LayerNorm(512)       │
+                    │  GELU()               │
+                    │  Dropout(0.3)         │
+                    │  Linear(512 → 256)    │
+                    │  LayerNorm(256)       │
+                    │  GELU()               │
+                    │  Dropout(0.2)         │
+                    │  Linear(256 → 6)      │
+                    └───────────┬───────────┘
+                                │
+                        ┌───────▼────────┐
+                        │ OUTPUT LOGITS  │
+                        │  (6 classes)   │
+                        └────────────────┘
+```
+
+**Optimizer**: AdamW with per-group learning rates
+- EfficientNet features[5]: lr=1e-5
+- EfficientNet features[6,7]: lr=3e-5
+- ConvNeXt stages[2,3] + layernorm: lr=3e-5
+- Fusion head: lr=1e-4
+- Loss: CrossEntropyLoss with label_smoothing=0.1
+- Early stopping patience: 7
+
+### ResNet-18 (Lightweight — 77%)
+
+```
+┌──────────────────────────────────┐
+│      INPUT IMAGE                 │
+│     (3, 128, 128)                │
+└───────────────┬──────────────────┘
+                │
+        ┌───────▼─────────┐
+        │   ResNet-18     │
+        │   Backbone      │
+        │                 │
+        │  Frozen except  │
+        │  layer3, layer4 │
+        └───────┬─────────┘
+                │  (512,)
+        ┌───────▼─────────────────────┐
+        │  Classification Head        │
+        │  Dropout(0.5)               │
+        │  Linear(512 → 256)          │
+        │  ReLU()                     │
+        │  Dropout(0.3)               │
+        │  Linear(256 → 6 classes)    │
+        └───────┬─────────────────────┘
+                │
+        ┌───────▼──────────┐
+        │  OUTPUT LOGITS   │
+        │  (6 classes)     │
+        └──────────────────┘
+```
+
+**Optimizer**: AdamW with per-group learning rates
+- layer3: lr=1e-5
+- layer4: lr=1e-5
+- fc head: lr=1e-4
+- Loss: CrossEntropyLoss
+- Early stopping patience: 7
+
+### YOLO v11m Integration
+
+```
+┌─────────────────────────────┐
+│   INPUT IMAGE               │
+│   imgsz=640, conf=0.05      │
+└──────────────┬──────────────┘
+               │
+       ┌───────▼────────┐
+       │  YOLO v11m     │
+       │  Inference     │
+       └───────┬────────┘
+               │
+    ┌──────────┴──────────┐
+    │                     │
+┌───▼───────┐      ┌──────▼──────┐
+│ Bboxes    │      │ Confidence  │
+│ (x1,y1,   │      │ Scores +    │
+│  x2,y2)   │      │ Class Label │
+└───┬───────┘      └──────┬──────┘
+    └──────────┬──────────┘
+               │
+       ┌───────▼────────┐
+       │ result.plot()  │
+       │ Save to disk   │
+       └────────────────┘
+```
+
+### Grad-CAM Pipeline (scripts/gradcam.py)
+
+```
+Image Path
+    │
+    ├─ ResNet mode:  target_layer = model.layer4[-1]
+    └─ Fusion mode:  target_layer = model.eff_features[-1]
+                     (FP16 → FP32 cast on CPU automatically)
+    │
+    ├─ Register forward hook  (_GradCAMHook)
+    ├─ Forward pass → score.backward()
+    ├─ acts [C,H,W]  ×  weights (mean of grads) → CAM [H,W]
+    ├─ ReLU → normalize → resize to original dims
+    └─ cv2.applyColorMap(COLORMAP_JET) → addWeighted overlay
+```
+
+### Data Pipeline (src/data/)
+
+```
+Raw Images (data/dataset/)
+    │
+    ├─ ingestion.py   → scan folders, build file list
+    ├─ preprocessing.py → validate / clean images
+    ├─ augmentation.py  → train/val transforms
+    │     ResNet:  Resize(128,128) + HFlip + Rotation(15°) + ColorJitter
+    │     Fusion:  Resize(260,260) + HFlip + Rotation(10°) + ColorJitter
+    └─ dataset.py   → ImageFolder DataLoaders
+                       (train 80% / val 20%, seed=42)
+```
+
+### Export & Deployment (src/export/)
+
+```
+Trained Checkpoints (checkpoints/)
+    │
+    ├─ conver_model.py         → FP32 → FP16 conversion
+    │                            788MB → 135MB (82.9% reduction)
+    └─ upload_to_huggingface.py → HfApi upload to:
+          junaid17/new-damagelens-resnet-classifier
+          junaid17/new-damagelens-fusion-fp16
+          junaid17/new-damagelens-yolo-detector
+```
+
+---
+
+## 📊 Model Performance
+
+### Fusion Model (High Accuracy — 84% Overall)
+
+**Classification Report:**
+
+![Fusion Classification Report](assets/fusion_classification_report.png)
+
+**Confusion Matrix:**
+
+![Fusion Confusion Matrix](assets/fusion_confusion_matrix.png)
+
+**Training Curves:**
+
+![Fusion Training Curves](assets/fusion_training_curves.png)
+
+---
+
+### ResNet-18 (Lightweight — 77% Overall)
+
+**Classification Report:**
+
+![ResNet Classification Report](assets/resnet_classification_report.png)
+
+**Confusion Matrix:**
+
+![ResNet Confusion Matrix](assets/resnet_confusion_matrix.png)
+
+**Training Curves:**
+
+![ResNet Training Curves](assets/resnet_training_curves.png)
+
+---
+
+### YOLO Detection Results
+
+![YOLO Detection Sample](assets/yolo_detection_sample.jpg)
+
+---
+
+## 🔁 CI Pipeline
+
+DamageLens uses **GitHub Actions** for continuous integration. Every push or pull request to `main`, `master`, or `dev` triggers the full test suite automatically.
+
+**CI Screenshot (GitHub Actions — All Tests Passing):**
+
+![CI Pipeline Passing](assets/ci_pipeline_passing.png)
+
+### What the pipeline tests:
+
+| Step | Test File | What it covers |
+|------|-----------|----------------|
+| Config | `test_config.py` | Paths, constants, class map |
+| Ingestion | `test_ingestion.py` | Dataset folder scanning |
+| Preprocessing | `test_preprocessing.py` | Image validation & cleaning |
+| Augmentation | `test_augmentation.py` | Transform pipelines |
+| Dataset | `test_dataset.py` | DataLoader creation |
+| ResNet Architecture | `test_resnet_model.py` | Model init & forward pass |
+| ResNet Training | `test_train_resnet.py` | Smoke test training loop |
+
+### Pipeline config (`.github/workflows/ci.yaml`):
+- Runs on: `ubuntu-latest`
+- Python: `3.10`
+- Triggers: push & PR to `main` / `master` / `dev`
+
 ---
-title: DamageLensAI
-emoji: 😻
-colorFrom: red
-colorTo: yellow
-sdk: docker
-pinned: false
+
+## 🚀 Setup & Installation
+
+### Prerequisites
+
+- Python 3.11+
+- CUDA 11.8+ (for GPU acceleration, optional but recommended)
+- 8GB+ RAM (16GB recommended for Fusion model)
+
+### Installation Steps
+
+```bash
+# Clone the repository
+git clone https://github.com/junaid17/damagelens.git
+cd DamageLens
+
+# Create virtual environment
+python -m venv myvenv
+source myvenv/bin/activate  # On Windows: myvenv\Scripts\activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Create required directories
+mkdir -p static/uploads static/results checkpoints assets
+```
+
+### Download Pre-trained Models
+
+Models are automatically downloaded from Hugging Face on first use:
+- `car-damage-classifier.pt` — ResNet-18 checkpoint
+- `best_fusion_model_fp16.pt` — Fusion model (FP16 optimized, 135MB)
+- `damage_detector.pt` — YOLO v11m model
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## 💻 Usage
+
+### Running the FastAPI Server
+
+```bash
+uvicorn app:app --reload --host 127.0.0.1 --port 8000
+```
+
+Open your browser at `http://127.0.0.1:8000`
+
+#### Quick Start:
+1. Upload a car image (JPG/PNG)
+2. Select analysis mode: **Fusion** (accurate) or **ResNet** (fast)
+3. Click "Run AI Analysis"
+4. View results in tabs:
+   - 📊 **Prediction**: Confidence scores and probabilities
+   - 👀 **Grad-CAM**: Visualize which regions influenced the prediction
+   - 🎯 **YOLO**: Damage bounding boxes with confidence
+
+### Python API Example
+
+```python
+import requests
+
+with open('car_image.jpg', 'rb') as f:
+    files = {'image': f}
+    resp = requests.post('http://127.0.0.1:8000/predict/resnet', files=files)
+    print(resp.json())
+
+with open('car_image.jpg', 'rb') as f:
+    files = {'image': f}
+    resp = requests.post('http://127.0.0.1:8000/predict/fusion', files=files)
+    print(resp.json())
+```
+
+---
+
+## 📡 API Documentation
+
+### `POST /predict/resnet`
+```
+Content-Type: multipart/form-data
+Body: image (File)
+
+Response:
+{
+  "status": "success",
+  "prediction": {
+    "Rear Normal": 0.47,
+    "Front Normal": 0.25,
+    ...
+  }
+}
+```
+
+### `POST /predict/fusion`
+```
+Content-Type: multipart/form-data
+Body: image (File)
+
+Response:
+{
+  "status": "success",
+  "prediction": {
+    "Rear Normal": 0.49,
+    "Front Normal": 0.35,
+    ...
+  }
+}
+```
+
+### `POST /predict?mode={resnet|fusion}` — Grad-CAM
+```
+Content-Type: multipart/form-data
+Body: file (File), mode (String)
+
+Response:
+{
+  "status": "success",
+  "mode": "fusion",
+  "original_image": "/static/uploads/{uuid}_input.jpg",
+  "selected_viz": "/static/results/{uuid}_fusion.jpg",
+  "resnet_viz": null,
+  "fusion_viz": "/static/results/{uuid}_fusion.jpg"
+}
+```
+
+### `POST /predict/yolo`
+```
+Content-Type: multipart/form-data
+Body: file (File)
+
+Response:
+{
+  "status": "success",
+  "original_image": "/static/uploads/{uuid}_input.jpg",
+  "yolo_image": "/static/results/{uuid}_yolo.jpg",
+  "detections": [
+    { "label": "damage", "confidence": 0.87, "box": [x1, y1, x2, y2] }
+  ],
+  "total_detections": 2,
+  "message": "Detections found"
+}
+```
+
+---
+
+## 🔧 Model Optimization
+
+### FP16 Conversion (Fusion Model)
+
+```
+Original Model (FP32):     788 MB
+Optimized Model (FP16):    135 MB
+───────────────────────────────────
+Compression Ratio:         82.9% reduction ✅
+Accuracy Loss:             < 1%            ⚠️
+Speed Improvement:         ~1.3x faster   ⚡
+```
+
+The system auto-detects FP16 checkpoints at load time:
+
+```python
+if first_tensor.dtype == torch.float16:
+    model = model.half()
+
+# Grad-CAM on CPU: FP16 → FP32 cast applied automatically
+if is_half:
+    model = model.float()
+```
+
+---
+
+## 📚 Dataset & Training
+
+### Data Constraints
+
+- **Total Samples**: ~1,800 images
+- **Train/Val Split**: 80/20 (seed=42)
+- **Classes**: 6 (F_Breakage, F_Crushed, F_Normal, R_Breakage, R_Crushed, R_Normal)
+- **YOLO subset**: ~100 annotated images (train/val split)
+
+### Data Augmentation
+
+| Transform | ResNet | Fusion |
+|-----------|--------|--------|
+| Resize | 128×128 | 260×260 |
+| RandomHorizontalFlip | ✅ | ✅ |
+| RandomRotation | ±15° | ±10° |
+| ColorJitter (b/c/s) | ±20% | ±15% |
+| ImageNet Normalize | ✅ | ✅ |
+
+### Training Configuration
+
+| Setting | ResNet | Fusion |
+|---------|--------|--------|
+| Backbone | ResNet-18 | EfficientNet-V2-S + ConvNeXt-Small |
+| Frozen layers | All except layer3, layer4 | All except features[5,6,7] / stages[2,3] |
+| Optimizer | AdamW | AdamW (per-group LR) |
+| Loss | CrossEntropyLoss | CrossEntropyLoss (label_smoothing=0.1) |
+| Early stopping | patience=7 | patience=7 |
+| Input size | 128×128 | 260×260 (EfficientNet) / 224×224 (ConvNeXt) |
+
+---
+
+## 🎨 Web UI Features
+
+- Dark mode glassmorphism design
+- Drag & drop image upload
+- Model selection dropdown (Fusion / ResNet)
+- Real-time confidence bar animation
+- Tab navigation: Prediction → Grad-CAM → YOLO
+- Scan line effect during processing
+- Plotly bar chart for class probabilities
+- Side-by-side original vs heatmap comparison
+
+---
+
+## 🔍 Grad-CAM Visualization
+
+Gradient-weighted Class Activation Mapping highlights which image regions most influenced the model's prediction.
+
+```
+Original Image    +    Grad-CAM Heatmap    =    Overlay
+                       Red   = High importance
+                       Blue  = Low importance
+```
+
+- ResNet: hooks into `layer4[-1]`
+- Fusion: hooks into `eff_features[-1]` (EfficientNet's last block)
+
+---
+
+## 📋 Directory Structure
+
+```
+DamageLens/
+├── app.py                              # FastAPI app + all endpoints
+├── index.html                          # Web UI
+├── requirements.txt
+├── README.md
+│
+├── .github/
+│   └── workflows/
+│       └── ci.yaml                     # GitHub Actions CI pipeline
+│
+├── assets/                             # ← Place README images here
+│   ├── fusion_classification_report.png
+│   ├── fusion_confusion_matrix.png
+│   ├── fusion_training_curves.png
+│   ├── resnet_classification_report.png
+│   ├── resnet_confusion_matrix.png
+│   ├── resnet_training_curves.png
+│   ├── yolo_detection_sample.png
+│   └── ci_pipeline_passing.png
+│
+├── scripts/
+│   ├── prediction_helper.py            # ResNet + Fusion model classes & inference
+│   ├── gradcam.py                      # Grad-CAM (ResNet + Fusion, CPU-optimized)
+│   ├── load_models.py                  # HF Hub download + model initialization
+│   └── yolo_predict.py                 # YOLO inference + bbox drawing
+│
+├── src/
+│   ├── config.py                       # Paths, hyperparams, class map
+│   ├── data/
+│   │   ├── ingestion.py                # Dataset folder scanning
+│   │   ├── preprocessing.py            # Image validation
+│   │   ├── augmentation.py             # Train/val transforms
+│   │   └── dataset.py                  # DataLoader creation
+│   ├── models/
+│   │   ├── resnet_model.py             # CarClassifierResNet
+│   │   └── fusion_model.py             # FusionClassifier
+│   ├── training/
+│   │   ├── trainer.py                  # Generic train loop (single + dual input)
+│   │   ├── train_resnet.py             # ResNet training entry point
+│   │   ├── train_fusion.py             # Fusion training entry point
+│   │   └── train_yolo.py               # YOLO fine-tuning
+│   └── export/
+│       ├── conver_model.py             # FP32 → FP16 conversion
+│       └── upload_to_huggingface.py    # HF Hub upload script
+│
+├── checkpoints/
+│   ├── best_resnet_model.pt
+│   ├── best_fusion_model_fp16.pt
+│   ├── damage_detector.pt
+│   └── yolo11m.pt
+│
+├── Notebooks/
+│   ├── Resnet18_fine_tuning_final.ipynb
+│   ├── EfficientNet_ConvNext_Fusion.ipynb
+│   └── damage_detector_yolo.ipynb
+│
+├── test/
+│   ├── test_config.py
+│   ├── test_ingestion.py
+│   ├── test_preprocessing.py
+│   ├── test_augmentation.py
+│   ├── test_dataset.py
+│   ├── test_resnet_model.py
+│   ├── test_fusion_model.py
+│   ├── test_train_resnet.py
+│   ├── test_train_fusion.py
+│   ├── test_train_yolo.py
+│   ├── test_model_conversion.py
+│   └── test_upload_to_huggingface.py
+│
+├── data/
+│   ├── dataset/                        # 6-class image folders
+│   │   ├── F_Breakage/
+│   │   ├── F_Crushed/
+│   │   ├── F_Normal/
+│   │   ├── R_Breakage/
+│   │   ├── R_Crushed/
+│   │   └── R_Normal/
+│   └── yolo/                           # YOLO annotated subset
+│       ├── train/images + labels/
+│       ├── val/images + labels/
+│       └── dataset_custom.yaml
+│
+└── static/
+    ├── uploads/                        # Temp uploaded images
+    └── results/                        # Generated Grad-CAM / YOLO outputs
+```
+
+---
+
+## ⚠️ Limitations & Known Issues
+
+### Data Constraints
+- **Limited Training Data**: ~1,800 samples — may show variance on edge cases
+- **Class Imbalance**: Rear Crushed class has fewer samples, affecting recall
+
+### Performance
+
+| Metric | Value | Note |
+|--------|-------|------|
+| ResNet Inference | ~500ms | Fast, lower accuracy |
+| Fusion Inference | 30-60s | Accurate, computationally heavy |
+| Cold Startup | 4-5 min | HF Hub download + model warmup |
+| GPU Memory | ~4GB | For Fusion model |
+| ResNet Accuracy | 77% | Lightweight trade-off |
+| Fusion Accuracy | 84% | Best accuracy |
+
+### Technical Limitations
+- Fusion accuracy is **7% higher** than ResNet (84% vs 77%)
+- YOLO model may miss small or partially occluded damage
+- Grad-CAM is for diagnostic/explainability purposes only
+- Batch processing not currently supported
+- FP16 Grad-CAM on CPU requires automatic FP32 cast (handled internally)

app.py

@@ -0,0 +1,236 @@
+import os
+import uuid
+import shutil
+import logging
+from contextlib import asynccontextmanager
+
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv
+
+from scripts.gradcam import get_resnet_gradcam, get_fusion_gradcam
+from scripts.yolo_predict import get_yolo_damage_boxes
+from scripts.load_models import initialize_models
+
+# ---------------- LOGGING ----------------
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+
+# ---------------- ENV ----------------
+load_dotenv()
+
+# ---------------- DIRECTORIES ----------------
+UPLOAD_DIR = "static/uploads"
+RESULT_DIR = "static/results"
+
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(RESULT_DIR, exist_ok=True)
+
+# ---------------- GLOBAL MODELS ----------------
+resnet_predictor = None
+fusion_predictor = None
+yolo_model = None
+
+CLASS_MAP = {
+    0: "Front Breakage",
+    1: "Front Crushed",
+    2: "Front Normal",
+    3: "Rear Breakage",
+    4: "Rear Crushed",
+    5: "Rear Normal"
+}
+
+# ---------------- FASTAPI STARTUP ----------------
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global resnet_predictor, fusion_predictor, yolo_model
+
+    logger.info("Loading models at startup...")
+
+    try:
+        resnet_predictor, fusion_predictor, yolo_model = initialize_models(CLASS_MAP)
+        logger.info("All models loaded successfully.")
+
+    except Exception as e:
+        logger.exception("Model loading failed.")
+        raise RuntimeError(str(e))
+
+    yield
+
+    logger.info("Application shutdown.")
+
+# ---------------- APP ----------------
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],   # restrict this in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# ---------------- HELPERS ----------------
+def validate_image(upload_file: UploadFile):
+    if not upload_file.content_type.startswith("image/"):
+        raise HTTPException(
+            status_code=400,
+            detail="Uploaded file must be an image."
+        )
+
+
+def save_upload(upload_file: UploadFile):
+    unique_id = str(uuid.uuid4())
+
+    filename = f"{unique_id}_input.jpg"
+    file_path = os.path.join(UPLOAD_DIR, filename)
+
+    with open(file_path, "wb") as buffer:
+        shutil.copyfileobj(upload_file.file, buffer)
+
+    return unique_id, filename, file_path
+
+# ---------------- ROUTES ----------------
+@app.get("/")
+def api_status():
+    return {"status": "API is running"}
+
+
+@app.post("/predict")
+async def predict_and_generate_cams(
+    file: UploadFile = File(...),
+    mode: str = "resnet"
+):
+    validate_image(file)
+
+    mode = mode.lower()
+
+    if mode not in {"resnet", "fusion"}:
+        raise HTTPException(
+            status_code=400,
+            detail="mode must be 'resnet' or 'fusion'"
+        )
+
+    try:
+        unique_id, input_filename, input_path = save_upload(file)
+
+        if mode == "resnet":
+            output_name = f"{unique_id}_resnet.jpg"
+            output_path = os.path.join(RESULT_DIR, output_name)
+
+            get_resnet_gradcam(
+                input_path,
+                resnet_predictor,
+                output_path
+            )
+
+            selected_viz = f"/static/results/{output_name}"
+
+            return {
+                "status": "success",
+                "mode": mode,
+                "original_image": f"/static/uploads/{input_filename}",
+                "selected_viz": selected_viz,
+                "resnet_viz": selected_viz,
+                "fusion_viz": None
+            }
+
+        output_name = f"{unique_id}_fusion.jpg"
+        output_path = os.path.join(RESULT_DIR, output_name)
+
+        get_fusion_gradcam(
+            input_path,
+            fusion_predictor,
+            output_path
+        )
+
+        selected_viz = f"/static/results/{output_name}"
+
+        return {
+            "status": "success",
+            "mode": mode,
+            "original_image": f"/static/uploads/{input_filename}",
+            "selected_viz": selected_viz,
+            "resnet_viz": None,
+            "fusion_viz": selected_viz
+        }
+
+    except Exception as e:
+        logger.exception("GradCAM generation failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/predict/resnet")
+async def resnet_prediction(image: UploadFile = File(...)):
+    validate_image(image)
+
+    try:
+        pil_image = Image.open(image.file).convert("RGB")
+
+        result = resnet_predictor.resnet_predict(pil_image)
+
+        return {
+            "status": "success",
+            "prediction": result
+        }
+
+    except Exception as e:
+        logger.exception("ResNet prediction failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/predict/fusion")
+async def fusion_prediction(image: UploadFile = File(...)):
+    validate_image(image)
+
+    try:
+        pil_image = Image.open(image.file).convert("RGB")
+
+        result = fusion_predictor.predict(pil_image)
+
+        return {
+            "status": "success",
+            "prediction": result
+        }
+
+    except Exception as e:
+        logger.exception("Fusion prediction failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/predict/yolo")
+async def yolo_detection(file: UploadFile = File(...)):
+    validate_image(file)
+
+    try:
+        unique_id, input_filename, input_path = save_upload(file)
+
+        output_name = f"{unique_id}_yolo.jpg"
+        output_path = os.path.join(RESULT_DIR, output_name)
+
+        result = get_yolo_damage_boxes(
+            input_path,
+            yolo_model,
+            output_path
+        )
+
+        return {
+            "status": "success",
+            "original_image": f"/static/uploads/{input_filename}",
+            "yolo_image": f"/static/results/{output_name}",
+            "detections": result["detections"],
+            "total_detections": result["total_detections"],
+            "message": result["message"]
+        }
+
+    except Exception as e:
+        logger.exception("YOLO detection failed.")
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -0,0 +1,15 @@
+torch
+torchvision
+transformers
+fastapi
+uvicorn
+dotenv
+matplotlib
+opencv-python
+python-multipart
+ultralytics
+plotly
+pandas
+scikit-learn
+seaborn
+huggingface_hub

scripts/gradcam.py

@@ -0,0 +1,167 @@
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# ------------------------------------------------------------------
+# Lightweight hook manager — CPU-only, no logging, direct capture
+# ------------------------------------------------------------------
+class _GradCAMHook:
+    __slots__ = ("activation", "gradient", "fwd_handle", "bwd_handle")
+
+    def __init__(self, target_layer):
+        self.activation = None
+        self.gradient = None
+        self.fwd_handle = target_layer.register_forward_hook(self._fwd_hook)
+        self.bwd_handle = None
+
+    def _fwd_hook(self, module, inp, out):
+        self.activation = out
+        # Tensor-level hook is lighter than full backward hook or retain_grad()
+        self.bwd_handle = out.register_hook(self._bwd_hook)
+
+    def _bwd_hook(self, grad):
+        self.gradient = grad
+
+    def remove(self):
+        self.fwd_handle.remove()
+        if self.bwd_handle is not None:
+            self.bwd_handle.remove()
+
+
+def _postprocess_cam(cam_tensor, original_img, output_path, alpha=0.5, beta=0.6):
+    """
+    CPU post-processing shared by both ResNet and Fusion.
+    cam_tensor: 2D torch tensor [H, W] on CPU, already ReLU'd
+    """
+    h, w = original_img.height, original_img.width
+
+    # Normalize on CPU (vectorized)
+    cam_min = cam_tensor.min()
+    cam_max = cam_tensor.max()
+    if cam_max > cam_min:
+        cam_tensor = (cam_tensor - cam_min) / (cam_max - cam_min)
+    else:
+        cam_tensor = torch.zeros_like(cam_tensor)
+
+    # Convert to numpy once, then resize with OpenCV (very fast on CPU)
+    cam_np = cam_tensor.numpy()
+    cam_np = cv2.resize(cam_np, (w, h), interpolation=cv2.INTER_LINEAR)
+
+    cam_np = np.uint8(255 * cam_np)
+    heatmap = cv2.applyColorMap(cam_np, cv2.COLORMAP_JET)
+
+    original_bgr = cv2.cvtColor(np.array(original_img), cv2.COLOR_RGB2BGR)
+    overlay = cv2.addWeighted(original_bgr, alpha, heatmap, beta, 0)
+    cv2.imwrite(output_path, overlay)
+
+
+# ------------------------------------------------------------------
+# Optimized ResNet Grad-CAM (CPU)
+# ------------------------------------------------------------------
+def get_resnet_gradcam(image_path, predictor, output_path):
+    logger.info("Starting ResNet Grad-CAM generation...")
+
+    model = predictor.model
+    model.eval()
+
+    target_layer = model.model.layer4[-1]
+    hook = _GradCAMHook(target_layer)
+
+    try:
+        original_img = Image.open(image_path).convert("RGB")
+        input_tensor = predictor.test_transforms(original_img).unsqueeze(0)
+
+        output = model(input_tensor)
+        score, pred_class_idx = output[0].max(dim=0)
+        pred_class_idx = pred_class_idx.item()
+
+        logger.info(f"Predicted class index: {pred_class_idx}")
+        score.backward()
+
+        if hook.activation is None or hook.gradient is None:
+            raise RuntimeError("Failed to capture activations or gradients.")
+
+        # ----- Vectorized Grad-CAM on CPU -----
+        acts = hook.activation[0].detach().float()     # [C, H, W]
+        grads = hook.gradient[0].detach().float()      # [C, H, W]
+
+        weights = grads.mean(dim=(1, 2), keepdim=True) # [C, 1, 1]
+        cam = (weights * acts).sum(dim=0)              # [H, W]
+        cam = F.relu(cam)
+
+        _postprocess_cam(cam, original_img, output_path, alpha=0.6, beta=0.4)
+
+        logger.info(f"ResNet Grad-CAM saved to: {output_path}")
+        return True
+
+    except Exception as e:
+        logger.exception("ResNet Grad-CAM generation failed.")
+        raise RuntimeError(f"ResNet Grad-CAM failed: {e}") from e
+
+    finally:
+        hook.remove()
+
+
+# ------------------------------------------------------------------
+# Optimized Fusion Grad-CAM (EfficientNet + ConvNeXt) (CPU)
+# ------------------------------------------------------------------
+def get_fusion_gradcam(image_path, predictor, output_path):
+    logger.info("Starting Fusion Grad-CAM generation...")
+
+    model = predictor.model
+    model.eval()
+
+    # FIX: PyTorch CPU does not support FP16 convolutions well.
+    # If the model is HalfTensor, cast it to FP32 for this pass.
+    is_half = next(model.parameters()).dtype == torch.float16
+    if is_half:
+        logger.info("FP16 model detected on CPU. Converting to FP32 for compatibility.")
+        model = model.float()
+
+    target_layer = model.eff_features[-1]
+    hook = _GradCAMHook(target_layer)
+
+    try:
+        original_img = Image.open(image_path).convert("RGB")
+
+        # CPU-only preprocessing (FloatTensor, no .to(device), no .half())
+        pixel_eff = predictor.eff_normalize(original_img).unsqueeze(0)
+        pixel_cnx = predictor.convnext_processor(
+            images=original_img, return_tensors="pt"
+        )["pixel_values"]
+
+        output = model(pixel_eff, pixel_cnx)
+        score, pred_class_idx = output[0].max(dim=0)
+        pred_class_idx = pred_class_idx.item()
+
+        logger.info(f"Predicted class index: {pred_class_idx}")
+        score.backward()
+
+        if hook.activation is None or hook.gradient is None:
+            raise RuntimeError("Failed to capture activations or gradients.")
+
+        # ----- Vectorized Grad-CAM on CPU -----
+        acts = hook.activation[0].detach().float()     # [C, H, W]
+        grads = hook.gradient[0].detach().float()        # [C, H, W]
+
+        weights = grads.mean(dim=(1, 2), keepdim=True) # [C, 1, 1]
+        cam = (weights * acts).sum(dim=0)              # [H, W]
+        cam = F.relu(cam)
+
+        _postprocess_cam(cam, original_img, output_path, alpha=0.5, beta=0.6)
+
+        logger.info(f"Fusion Grad-CAM saved to: {output_path}")
+        return True
+
+    except Exception as e:
+        logger.exception("Fusion Grad-CAM generation failed.")
+        raise RuntimeError(f"Fusion Grad-CAM failed: {e}") from e
+
+    finally:
+        hook.remove()

scripts/load_models.py

@@ -0,0 +1,91 @@
+import logging
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+from .prediction_helper import (
+    ResnetCarDamagePredictor,
+    FusionCarDamagePredictor,
+)
+
+logger = logging.getLogger(__name__)
+
+MODEL_CONFIG = {
+    "resnet": {
+        "repo_id": "junaid17/car-damage-classifier",
+        "filename": "car-damage-classifier.pt",
+    },
+    "fusion": {
+        "repo_id": "junaid17/best_fusion_model_fp16",
+        "filename": "best_fusion_model_fp16.pt",
+    },
+    "yolo": {
+        "repo_id": "junaid17/Yolo_Model",
+        "filename": "damage_detector.pt",
+    },
+}
+
+
+def get_checkpoint_path(model_key: str) -> Path:
+    if model_key not in MODEL_CONFIG:
+        raise ValueError(f"Unknown model key: {model_key}")
+
+    config = MODEL_CONFIG[model_key]
+
+    try:
+        logger.info(f"Fetching {model_key} model from Hugging Face Hub...")
+        logger.info(f"Repo: {config['repo_id']}")
+        logger.info(f"File: {config['filename']}")
+
+        local_path = hf_hub_download(
+            repo_id=config["repo_id"],
+            filename=config["filename"],
+        )
+
+        logger.info(f"{model_key} model available at: {local_path}")
+
+        return Path(local_path)
+
+    except Exception as e:
+        logger.exception(f"Failed to fetch {model_key} model.")
+        raise RuntimeError(f"Failed to load {model_key} checkpoint: {str(e)}")
+
+
+class ModelLoader:
+    def __init__(self):
+        logger.info("Initializing ModelLoader...")
+
+    def get_model_path(self, model_key: str) -> Path:
+        return get_checkpoint_path(model_key)
+
+
+def initialize_models(class_map):
+    logger.info("Starting model initialization...")
+
+    try:
+        resnet_path = get_checkpoint_path("resnet")
+        fusion_path = get_checkpoint_path("fusion")
+        yolo_path = get_checkpoint_path("yolo")
+
+        logger.info("Initializing ResNet predictor...")
+        resnet_predictor = ResnetCarDamagePredictor(
+            checkpoint_path=resnet_path,
+            class_map=class_map
+        )
+
+        logger.info("Initializing Fusion predictor...")
+        fusion_predictor = FusionCarDamagePredictor(
+            checkpoint_path=fusion_path,
+            class_map=class_map
+        )
+
+        logger.info("Initializing YOLO model...")
+        yolo_model = YOLO(str(yolo_path))
+
+        logger.info("All models initialized successfully.")
+
+        return resnet_predictor, fusion_predictor, yolo_model
+
+    except Exception as e:
+        logger.exception("Model initialization failed.")
+        raise RuntimeError(f"Model initialization failed: {str(e)}")

scripts/prediction_helper.py

@@ -0,0 +1,313 @@
+import os
+import logging
+import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from PIL import Image, UnidentifiedImageError
+from transformers import ConvNextModel, ConvNextImageProcessor
+
+# ---------------- LOGGING SETUP ----------------
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------- RESNET MODEL ----------------
+class Car_Classifier_Resnet(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+
+        logger.info("Initializing ResNet18 architecture...")
+
+        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)
+        )
+
+        logger.info("ResNet architecture initialized successfully.")
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class ResnetCarDamagePredictor:
+    def __init__(self, checkpoint_path, class_map):
+        logger.info("Initializing ResNet predictor...")
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.class_map = class_map
+
+        logger.info(f"Using device for ResNet: {self.device}")
+
+        self.test_transforms = transforms.Compose([
+            transforms.Resize((128, 128)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                [0.485, 0.456, 0.406],
+                [0.229, 0.224, 0.225]
+            )
+        ])
+
+        try:
+            self.model = Car_Classifier_Resnet(num_classes=len(class_map))
+
+            logger.info(f"Loading ResNet checkpoint from: {checkpoint_path}")
+
+            checkpoint = torch.load(checkpoint_path, map_location=self.device)
+            state_dict = checkpoint.get("model_state_dict", checkpoint)
+
+            self.model.load_state_dict(state_dict)
+            self.model.to(self.device)
+            self.model.eval()
+
+            logger.info("ResNet model loaded successfully.")
+
+        except Exception as e:
+            logger.exception("Failed to load ResNet model.")
+            raise RuntimeError(f"Failed to load ResNet model: {str(e)}")
+
+    def resnet_predict(self, image_input):
+        logger.info("Starting ResNet prediction...")
+
+        try:
+            if isinstance(image_input, str):
+                logger.info(f"Loading image from file path: {image_input}")
+                image = Image.open(image_input).convert("RGB")
+
+            elif isinstance(image_input, Image.Image):
+                logger.info("Using PIL image input.")
+                image = image_input.convert("RGB")
+
+            else:
+                raise TypeError("image_input must be a file path or PIL.Image")
+
+            image = self.test_transforms(image)
+            image = image.unsqueeze(0).to(self.device)
+
+            with torch.no_grad():
+                outputs = self.model(image)
+
+            probs = torch.nn.functional.softmax(outputs, dim=1)[0]
+
+            class_probs = {
+                self.class_map[i]: float(probs[i].item())
+                for i in range(len(self.class_map))
+            }
+
+            sorted_probs = dict(
+                sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+            )
+
+            logger.info("ResNet prediction completed successfully.")
+
+            return sorted_probs
+
+        except UnidentifiedImageError:
+            logger.error("Invalid image file provided to ResNet predictor.")
+            raise ValueError("Invalid image file provided")
+
+        except Exception as e:
+            logger.exception("ResNet prediction failed.")
+            raise RuntimeError(f"ResNet prediction failed: {str(e)}")
+
+
+# ---------------- FUSION MODEL ----------------
+class FusionClassifier(nn.Module):
+    def __init__(self, num_classes, convnext_model_name="facebook/convnext-small-224"):
+        super().__init__()
+
+        logger.info("Initializing Fusion model architecture...")
+
+        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
+
+        logger.info("Loading ConvNeXt 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
+
+        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)
+        )
+
+        logger.info("Fusion architecture 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)
+
+        return self.fusion_head(fused)
+
+
+class FusionCarDamagePredictor:
+    def __init__(self, checkpoint_path, class_map, convnext_model_name="facebook/convnext-small-224"):
+        logger.info("Initializing Fusion predictor...")
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.class_map = class_map
+
+        logger.info(f"Using device for Fusion: {self.device}")
+
+        self.eff_normalize = transforms.Compose([
+            transforms.Resize((260, 260)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                [0.485, 0.456, 0.406],
+                [0.229, 0.224, 0.225]
+            )
+        ])
+
+        logger.info("Loading ConvNeXt image processor...")
+        self.convnext_processor = ConvNextImageProcessor.from_pretrained(
+            convnext_model_name
+        )
+
+        try:
+            self.model = FusionClassifier(
+                num_classes=len(class_map),
+                convnext_model_name=convnext_model_name
+            )
+
+            logger.info(f"Loading Fusion checkpoint from: {checkpoint_path}")
+
+            checkpoint = torch.load(checkpoint_path, map_location=self.device)
+            state_dict = checkpoint.get("model_state_dict", checkpoint)
+
+            first_tensor = next(iter(state_dict.values()))
+
+            if first_tensor.dtype == torch.float16:
+                logger.info("FP16 checkpoint detected. Converting model to half precision.")
+                self.model = self.model.half()
+
+            self.model.load_state_dict(state_dict)
+            self.model.to(self.device)
+            self.model.eval()
+
+            logger.info("Fusion model loaded successfully.")
+
+        except Exception as e:
+            logger.exception("Failed to load Fusion model.")
+            raise RuntimeError(f"Failed to load Fusion model: {str(e)}")
+
+    def predict(self, image_input):
+        logger.info("Starting Fusion prediction...")
+
+        try:
+            if isinstance(image_input, str):
+                logger.info(f"Loading image from file path: {image_input}")
+                image = Image.open(image_input).convert("RGB")
+
+            elif isinstance(image_input, Image.Image):
+                logger.info("Using PIL image input.")
+                image = image_input.convert("RGB")
+
+            else:
+                raise TypeError("image_input must be a file path or PIL.Image")
+
+            pixel_eff = self.eff_normalize(image)
+            pixel_eff = pixel_eff.unsqueeze(0).to(self.device)
+
+            inputs_cnx = self.convnext_processor(
+                images=image,
+                return_tensors="pt"
+            )
+
+            pixel_cnx = inputs_cnx["pixel_values"].to(self.device)
+
+            if next(self.model.parameters()).dtype == torch.float16:
+                logger.info("Converting input tensors to FP16.")
+                pixel_eff = pixel_eff.half()
+                pixel_cnx = pixel_cnx.half()
+
+            with torch.no_grad():
+                logits = self.model(pixel_eff, pixel_cnx)
+                probs = torch.nn.functional.softmax(logits, dim=1)[0]
+
+            class_probs = {
+                self.class_map[i]: float(probs[i].item())
+                for i in range(len(self.class_map))
+            }
+
+            sorted_probs = dict(
+                sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+            )
+
+            logger.info("Fusion prediction completed successfully.")
+
+            return sorted_probs
+
+        except UnidentifiedImageError:
+            logger.error("Invalid image file provided to Fusion predictor.")
+            raise ValueError("Invalid image file provided")
+
+        except Exception as e:
+            logger.exception("Fusion prediction failed.")
+            raise RuntimeError(f"Fusion prediction failed: {str(e)}")

scripts/yolo_predict.py

@@ -0,0 +1,63 @@
+import cv2
+import logging
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+def get_yolo_damage_boxes(image_path, yolo_model, output_path):
+    logger.info("Starting YOLO damage detection...")
+
+    try:
+        image = Image.open(image_path).convert("RGB")
+
+        results = yolo_model.predict(
+            source=image,
+            conf=0.05,
+            imgsz=640,
+            verbose=False
+        )
+
+        result = results[0]
+        boxes = result.boxes
+        detections = []
+
+        if boxes is not None and len(boxes) > 0:
+            logger.info(f"{len(boxes)} detections found.")
+
+            for box in boxes:
+                conf = float(box.conf[0])
+                cls_id = int(box.cls[0])
+
+                label = yolo_model.names[cls_id]
+
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+
+                detections.append({
+                    "label": label,
+                    "confidence": round(conf, 4),
+                    "box": [x1, y1, x2, y2]
+                })
+
+        else:
+            logger.info("No detections found.")
+
+        plotted = result.plot()
+
+        cv2.imwrite(output_path, plotted)
+
+        logger.info(f"YOLO output saved to: {output_path}")
+
+        return {
+            "detections": detections,
+            "total_detections": len(detections),
+            "message": (
+                "No damage detected"
+                if len(detections) == 0
+                else "Detections found"
+            )
+        }
+
+    except Exception as e:
+        logger.exception("YOLO detection failed.")
+        raise RuntimeError(f"YOLO failed: {str(e)}")

src/config.py

@@ -0,0 +1,60 @@
+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 = 128
+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)
+
+# ---------------- HUGGING FACE ----------------
+HF_USERNAME = "junaid17"
+
+HF_RESNET_REPO = "new-car-damage-classifier"
+HF_FUSION_REPO = "new-best-fusion-model-fp16"
+HF_YOLO_REPO = "new-Yolo-Model"

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/export/conver_model.py

@@ -0,0 +1,68 @@
+import os
+import logging
+import torch
+
+from src.config import DEVICE, NUM_CLASSES, CHECKPOINT_DIR
+from src.models.fusion_model import FusionClassifier
+
+logger = logging.getLogger(__name__)
+
+INPUT_CHECKPOINT = CHECKPOINT_DIR / "best_fusion_model.pt"
+OUTPUT_CHECKPOINT = CHECKPOINT_DIR / "best_fusion_model_fp16.pt"
+
+
+def convert_fusion_to_fp16():
+    logger.info("Initializing Fusion model for FP16 conversion...")
+
+    if not INPUT_CHECKPOINT.exists():
+        raise FileNotFoundError(
+            f"Fusion checkpoint not found: {INPUT_CHECKPOINT}"
+        )
+
+    model = FusionClassifier(
+        num_classes=NUM_CLASSES
+    ).to(DEVICE)
+
+    logger.info(f"Loading checkpoint from: {INPUT_CHECKPOINT}")
+
+    checkpoint = torch.load(
+        INPUT_CHECKPOINT,
+        map_location=DEVICE
+    )
+
+    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+        model.load_state_dict(checkpoint["model_state_dict"])
+    else:
+        model.load_state_dict(checkpoint)
+
+    logger.info("Model weights loaded successfully.")
+
+    model.eval()
+
+    logger.info("Converting model to FP16...")
+
+    model = model.half()
+
+    torch.save(
+        model.state_dict(),
+        OUTPUT_CHECKPOINT
+    )
+
+    size_mb = os.path.getsize(OUTPUT_CHECKPOINT) / (1024 * 1024)
+
+    logger.info(f"FP16 model saved at: {OUTPUT_CHECKPOINT}")
+    logger.info(f"FP16 model size: {size_mb:.2f} MB")
+
+    return OUTPUT_CHECKPOINT
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    fp16_path = convert_fusion_to_fp16()
+
+    print("\nFusion FP16 conversion completed successfully.")
+    print(f"Saved model: {fp16_path}")

src/export/upload_to_huggingface.py

@@ -0,0 +1,90 @@
+import os
+import logging
+from dotenv import load_dotenv
+from huggingface_hub import HfApi
+
+from src.config import CHECKPOINT_DIR
+
+load_dotenv()
+
+logger = logging.getLogger(__name__)
+
+HF_USERNAME = "junaid17"
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+if not HF_TOKEN:
+    raise ValueError(
+        "HF_TOKEN not found in .env file."
+    )
+
+MODELS = {
+    "new-damagelens-resnet-classifier": {
+        "path": CHECKPOINT_DIR / "best_resnet_model.pt",
+        "filename": "new_best_resnet_model.pt"
+    },
+
+    "new-damagelens-fusion-fp16": {
+        "path": CHECKPOINT_DIR / "best_fusion_model_fp16.pt",
+        "filename": "new_best_fusion_model_fp16.pt"
+    },
+
+    "new-damagelens-yolo-detector": {
+        "path": CHECKPOINT_DIR / "damage_detector.pt",
+        "filename": "new_damage_detector.pt"
+    }
+}
+
+
+def upload_model(api, repo_name, file_path, filename):
+    if not file_path.exists():
+        raise FileNotFoundError(
+            f"Model file not found: {file_path}"
+        )
+
+    repo_id = f"{HF_USERNAME}/{repo_name}"
+
+    logger.info(f"Creating repo: {repo_id}")
+
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type="model",
+        exist_ok=True
+    )
+
+    logger.info(f"Uploading {filename} to {repo_id}")
+
+    api.upload_file(
+        path_or_fileobj=str(file_path),
+        path_in_repo=filename,
+        repo_id=repo_id,
+        repo_type="model"
+    )
+
+    logger.info(f"Upload completed: {repo_id}")
+
+
+def upload_all_models():
+    logger.info("Starting Hugging Face model uploads...")
+
+    api = HfApi(token=HF_TOKEN)
+
+    for repo_name, model_info in MODELS.items():
+        upload_model(
+            api=api,
+            repo_name=repo_name,
+            file_path=model_info["path"],
+            filename=model_info["filename"]
+        )
+
+    logger.info("All model uploads completed successfully.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    upload_all_models()
+
+    print("\nAll models uploaded successfully.")

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_model_conversion.py

@@ -0,0 +1,39 @@
+import logging
+import os
+
+from src.export.conver_model import convert_fusion_to_fp16
+from src.config import CHECKPOINT_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def test_model_conversion():
+    logger.info("Testing fusion FP16 conversion...")
+
+    input_checkpoint = CHECKPOINT_DIR / "best_fusion_model.pt"
+
+    assert input_checkpoint.exists(), \
+        f"Missing checkpoint: {input_checkpoint}"
+
+    output_path = convert_fusion_to_fp16()
+
+    assert output_path.exists(), \
+        "FP16 model was not created"
+
+    size_mb = os.path.getsize(output_path) / (1024 * 1024)
+
+    assert size_mb > 0, \
+        "Generated FP16 model is empty"
+
+    logger.info("Model conversion test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_model_conversion()
+
+    print("Model conversion 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.")

test/test_upload_to_huggingface.py

@@ -0,0 +1,53 @@
+import logging
+import os
+from dotenv import load_dotenv
+from huggingface_hub import HfApi
+
+from src.export.upload_to_huggingface import MODELS
+
+logger = logging.getLogger(__name__)
+
+
+def test_huggingface_upload_setup():
+    logger.info("Testing Hugging Face upload setup...")
+
+    load_dotenv()
+
+    hf_token = os.getenv("HF_TOKEN")
+
+    assert hf_token is not None, \
+        "HF_TOKEN missing in .env"
+
+    assert hf_token.startswith("hf_"), \
+        "Invalid Hugging Face token format"
+
+    api = HfApi(token=hf_token)
+
+    assert api is not None, \
+        "Failed to initialize Hugging Face API"
+
+    for repo_name, model_info in MODELS.items():
+        file_path = model_info["path"]
+        filename = model_info["filename"]
+
+        assert file_path.exists(), \
+            f"Missing model file: {file_path}"
+
+        assert filename.endswith(".pt"), \
+            f"Invalid model filename: {filename}"
+
+        assert repo_name.startswith("new-"), \
+            f"Repo naming invalid: {repo_name}"
+
+    logger.info("Hugging Face upload setup test passed.")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    test_huggingface_upload_setup()
+
+    print("Hugging Face upload test completed successfully.")