Upload 6 files

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"]

requirements.txt

@@ -0,0 +1,23 @@
+torch
+torchvision
+transformers
+fastapi
+uvicorn
+dotenv
+matplotlib
+opencv-python
+python-multipart
+ultralytics
+plotly
+pandas
+scikit-learn
+seaborn
+huggingface_hub
+onnx 
+onnxruntime 
+pillow  
+onnxscript
+huggingface_hub
+python-dotenv
+onnxconverter-common
+streamlit

scripts/gradcam.py

@@ -0,0 +1,100 @@
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+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
+        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.6, beta=0.4):
+    h, w = original_img.height, original_img.width
+
+    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)
+
+    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)
+
+
+def get_resnet_gradcam(image_path, predictor, output_path):
+    """Generates Grad-CAM visual heatmap using the PyTorch ResNet classifier."""
+    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).to(predictor.device)
+
+        output = model(input_tensor)
+        score, pred_class_idx = output[0].max(dim=0)
+
+        logger.info(f"Predicted class index for Grad-CAM: {pred_class_idx.item()}")
+        score.backward()
+
+        if hook.activation is None or hook.gradient is None:
+            raise RuntimeError("Failed to capture activations or gradients.")
+
+        acts = hook.activation[0].detach().cpu().float()
+        grads = hook.gradient[0].detach().cpu().float()
+
+        weights = grads.mean(dim=(1, 2), keepdim=True)
+        cam = (weights * acts).sum(dim=0)
+        cam = F.relu(cam)
+
+        _postprocess_cam(cam, original_img, output_path)
+
+        logger.info(f"ResNet Grad-CAM heatmap 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()
+
+
+def get_fusion_gradcam(image_path, predictor, output_path):
+    """Fallback handler mapping to ResNet Grad-CAM (since single Grad-CAM architecture is used)."""
+    logger.info("Fusion Grad-CAM requested. Directing to ResNet Grad-CAM engine...")
+    return get_resnet_gradcam(image_path, predictor, output_path)

scripts/load_models.py

@@ -0,0 +1,104 @@
+import logging
+from pathlib import Path
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+from .prediction_helper import ResnetCarDamagePredictor
+
+logger = logging.getLogger(__name__)
+
+MODEL_CONFIG = {
+    "resnet_onnx": {
+        "repo_id": "junaid17/car-damage-classifier",
+        "filename": "car-damage-classifier.onnx",
+    },
+    "fusion_onnx": {
+        "repo_id": "junaid17/best_fusion_model_fp16",
+        "filename": "fusion_model.onnx",
+    },
+    "yolo_onnx": {
+        "repo_id": "junaid17/Yolo_Model",
+        "filename": "damage_detector.onnx",
+    },
+    "resnet_pt": {
+        "repo_id": "junaid17/car-damage-classifier",
+        "filename": "car-damage-classifier.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']} | File: {config['filename']}")
+
+        local_path = hf_hub_download(
+            repo_id=config["repo_id"],
+            filename=config["filename"],
+        )
+
+        logger.info(f"{model_key} model downloaded to: {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 pipeline...")
+
+    try:
+        # 1. Download/Fetch all 4 model file paths
+        resnet_onnx_path = get_checkpoint_path("resnet_onnx")
+        fusion_onnx_path = get_checkpoint_path("fusion_onnx")
+        yolo_onnx_path   = get_checkpoint_path("yolo_onnx")
+        resnet_pt_path   = get_checkpoint_path("resnet_pt")
+
+        # Define ONNX Execution Providers (GPU if available, fallback to CPU)
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+
+        # 2. Initialize ONNX Runtime Sessions for Classifier Models
+        logger.info("Initializing ResNet ONNX Session...")
+        resnet_onnx_session = ort.InferenceSession(str(resnet_onnx_path), providers=providers)
+
+        logger.info("Initializing Fusion ONNX Session...")
+        fusion_onnx_session = ort.InferenceSession(str(fusion_onnx_path), providers=providers)
+
+        # 3. Initialize YOLO ONNX Model via Ultralytics (enables ONNX runtime with .predict API)
+        logger.info("Initializing YOLO ONNX Model via Ultralytics...")
+        yolo_onnx_model = YOLO(str(yolo_onnx_path), task="detect")
+
+        # 4. Initialize PyTorch ResNet Predictor (specifically reserved for Grad-CAM)
+        logger.info("Initializing PyTorch ResNet model for Grad-CAM...")
+        resnet_gradcam_predictor = ResnetCarDamagePredictor(
+            checkpoint_path=resnet_pt_path,
+            class_map=class_map
+        )
+
+        logger.info("All 4 models (3 ONNX + 1 PyTorch Grad-CAM) initialized successfully.")
+
+        return {
+            "resnet_onnx": resnet_onnx_session,
+            "fusion_onnx": fusion_onnx_session,
+            "yolo_onnx": yolo_onnx_model,
+            "resnet_pt": resnet_gradcam_predictor,
+        }
+
+    except Exception as e:
+        logger.exception("Model initialization failed.")
+        raise RuntimeError(f"Model initialization failed: {str(e)}")

scripts/prediction_helper.py

@@ -0,0 +1,212 @@
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from PIL import Image, UnidentifiedImageError
+from transformers import ConvNextImageProcessor
+import onnxruntime as ort
+
+# ---------------- LOGGING SETUP ----------------
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+
+
+# =====================================================================
+# 1. PYTORCH RESNET MODEL & PREDICTOR (RETAINED EXCLUSIVELY FOR GRAD-CAM)
+# =====================================================================
+class Car_Classifier_Resnet(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        logger.info("Initializing ResNet18 PyTorch architecture for Grad-CAM...")
+
+        self.model = models.resnet18(weights="DEFAULT")
+
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        for param in self.model.layer3.parameters():
+            param.requires_grad = True
+
+        for param in self.model.layer4.parameters():
+            param.requires_grad = True
+
+        self.model.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(self.model.fc.in_features, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class ResnetCarDamagePredictor:
+    """PyTorch Predictor wrapper used strictly by Grad-CAM for backpropagation."""
+    def __init__(self, checkpoint_path, class_map):
+        logger.info("Initializing PyTorch ResNet predictor for Grad-CAM...")
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.class_map = class_map
+
+        self.test_transforms = transforms.Compose([
+            transforms.Resize((224, 224)),
+            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))
+            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("PyTorch ResNet model loaded successfully for Grad-CAM.")
+
+        except Exception as e:
+            logger.exception("Failed to load PyTorch ResNet model.")
+            raise RuntimeError(f"Failed to load ResNet PyTorch model: {str(e)}")
+
+
+# =====================================================================
+# 2. ONNX RESNET PREDICTOR (FAST INFERENCE)
+# =====================================================================
+class ResnetONNXPredictor:
+    def __init__(self, session: ort.InferenceSession, class_map: dict):
+        logger.info("Initializing ResNet ONNX Predictor...")
+        self.session = session
+        self.class_map = class_map
+        self.input_name = self.session.get_inputs()[0].name
+
+        self.transforms = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                [0.485, 0.456, 0.406],
+                [0.229, 0.224, 0.225]
+            )
+        ])
+
+    def predict(self, image_input):
+        logger.info("Starting ResNet ONNX prediction...")
+        try:
+            if isinstance(image_input, str):
+                image = Image.open(image_input).convert("RGB")
+            elif isinstance(image_input, Image.Image):
+                image = image_input.convert("RGB")
+            else:
+                raise TypeError("image_input must be a file path or PIL.Image")
+
+            # Preprocess to [1, 3, 224, 224] NumPy array
+            input_tensor = self.transforms(image).unsqueeze(0).numpy()
+
+            # Execute ONNX forward pass
+            outputs = self.session.run(None, {self.input_name: input_tensor})
+            logits = outputs[0][0]
+
+            # Softmax calculation
+            exp_logits = np.exp(logits - np.max(logits))
+            probs = exp_logits / np.sum(exp_logits)
+
+            class_probs = {
+                self.class_map[i]: float(probs[i])
+                for i in range(len(self.class_map))
+            }
+
+            sorted_probs = dict(
+                sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+            )
+
+            return sorted_probs
+
+        except UnidentifiedImageError:
+            logger.error("Invalid image file provided to ResNet ONNX predictor.")
+            raise ValueError("Invalid image file provided")
+        except Exception as e:
+            logger.exception("ResNet ONNX prediction failed.")
+            raise RuntimeError(f"ResNet ONNX prediction failed: {str(e)}")
+
+
+# =====================================================================
+# 3. ONNX FUSION PREDICTOR (EFFICIENTNET + CONVNEXT)
+# =====================================================================
+class FusionONNXPredictor:
+    def __init__(self, session: ort.InferenceSession, class_map: dict, convnext_model_name="facebook/convnext-small-224"):
+        logger.info("Initializing Fusion ONNX Predictor...")
+        self.session = session
+        self.class_map = class_map
+
+        # Dynamically retrieve ONNX input names (EffNet & ConvNeXt)
+        self.input_eff_name = self.session.get_inputs()[0].name
+        self.input_cnx_name = self.session.get_inputs()[1].name
+
+        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]
+            )
+        ])
+
+        self.convnext_processor = ConvNextImageProcessor.from_pretrained(convnext_model_name)
+
+    def predict(self, image_input):
+        logger.info("Starting Fusion ONNX prediction...")
+        try:
+            if isinstance(image_input, str):
+                image = Image.open(image_input).convert("RGB")
+            elif isinstance(image_input, Image.Image):
+                image = image_input.convert("RGB")
+            else:
+                raise TypeError("image_input must be a file path or PIL.Image")
+
+            # EfficientNet preprocessing -> [1, 3, 260, 260]
+            pixel_eff_np = self.eff_normalize(image).unsqueeze(0).numpy()
+
+            # ConvNeXt preprocessing -> [1, 3, 224, 224]
+            inputs_cnx = self.convnext_processor(images=image, return_tensors="pt")
+            pixel_cnx_np = inputs_cnx["pixel_values"].numpy()
+
+            # Execute dual-input ONNX forward pass
+            onnx_inputs = {
+                self.input_eff_name: pixel_eff_np,
+                self.input_cnx_name: pixel_cnx_np
+            }
+
+            outputs = self.session.run(None, onnx_inputs)
+            logits = outputs[0][0]
+
+            # Softmax calculation
+            exp_logits = np.exp(logits - np.max(logits))
+            probs = exp_logits / np.sum(exp_logits)
+
+            class_probs = {
+                self.class_map[i]: float(probs[i])
+                for i in range(len(self.class_map))
+            }
+
+            sorted_probs = dict(
+                sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+            )
+
+            return sorted_probs
+
+        except UnidentifiedImageError:
+            logger.error("Invalid image file provided to Fusion ONNX predictor.")
+            raise ValueError("Invalid image file provided")
+        except Exception as e:
+            logger.exception("Fusion ONNX prediction failed.")
+            raise RuntimeError(f"Fusion ONNX prediction failed: {str(e)}")

scripts/yolo_predict.py

@@ -0,0 +1,64 @@
+import cv2
+import logging
+from PIL import Image
+from ultralytics import YOLO
+
+logger = logging.getLogger(__name__)
+
+
+def get_yolo_damage_boxes(image_path, yolo_model: YOLO, output_path: str):
+    """Runs object detection using the YOLO ONNX model."""
+    logger.info("Starting YOLO ONNX damage detection...")
+
+    try:
+        image = Image.open(image_path).convert("RGB")
+
+        # Run ONNX inference through Ultralytics wrapper
+        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)} damage 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 damage detections found.")
+
+        # Save bounding box visual overlay
+        plotted = result.plot()
+        cv2.imwrite(output_path, plotted)
+
+        logger.info(f"YOLO ONNX 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 ONNX detection failed.")
+        raise RuntimeError(f"YOLO detection failed: {str(e)}")