Update predict_helper.py

predict_helper.py

@@ -1,3 +1,6 @@
+import base64
+import io
+import numpy as np
 from PIL import Image
 from torchvision import transforms
 from ultralytics import YOLO
@@ -12,30 +15,23 @@ yolo_model = YOLO("artifacts/damage_detector.pt")
 class Car_Classifier_Resnet(nn.Module):
     def __init__(self, num_classes):
         super().__init__()
-
         self.model = models.resnet18(weights="DEFAULT")
-
         for param in self.model.parameters():
             param.requires_grad = False
-
         for param in self.model.layer4.parameters():
             param.requires_grad = True
-
         for module in self.model.modules():
             if isinstance(module, nn.BatchNorm2d):
                 for param in module.parameters():
                     param.requires_grad = True
-
         self.model.fc = nn.Sequential(
             nn.Dropout(0.4),
             nn.Linear(self.model.fc.in_features, num_classes)
         )
 
-
     def forward(self, x):
         return self.model(x)
 
-
 class_names = [
     "F_Breakage",
     "F_Crushed",
@@ -46,7 +42,6 @@ class_names = [
 ]
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 clf_model = Car_Classifier_Resnet(num_classes=6).to(device)
 clf_model.load_state_dict(
     torch.load("artifacts/Damage_Classifier_Resnet_18.pth", map_location=device)
@@ -62,39 +57,51 @@ transform = transforms.Compose([
     )
 ])
 
-# here
-
 def predict_damage(image: Image.Image):
     image = image.convert("RGB")
+    
+    # -------- 1. CLASSIFICATION (ResNet) --------
+    # Run classification first as requested
+    img_tensor = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        out = clf_model(img_tensor)
+        probs = torch.softmax(out, dim=1)
+        conf, idx = torch.max(probs, dim=1)
+    
+    damage_type = class_names[idx.item()]
+    confidence_score = round(conf.item(), 4)
 
-    # -------- YOLO --------
+    # -------- 2. YOLO DETECTION --------
     yolo_results = yolo_model.predict(
         source=image,
         conf=0.05,
         imgsz=640,
         verbose=False
     )
-
-    bboxes = []
-    if yolo_results[0].boxes is not None:
-        for box in yolo_results[0].boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
-            conf = float(box.conf[0])
-            bboxes.append({
-                "bbox": [x1, y1, x2, y2],
-                "confidence": round(conf, 4)
-            })
-
-    # -------- CLASSIFICATION --------
-    img_tensor = transform(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        out = clf_model(img_tensor)
-        probs = torch.softmax(out, dim=1)
-        conf, idx = torch.max(probs, dim=1)
+    
+    result = yolo_results[0]
+    
+    # Check if any boxes were detected
+    damage_detected = result.boxes is not None and len(result.boxes) > 0
+
+    # Generate the image with bounding boxes drawn
+    # plot() returns a numpy array in BGR format (OpenCV style)
+    plotted_image_bgr = result.plot()
+    
+    # Convert BGR to RGB
+    plotted_image_rgb = plotted_image_bgr[..., ::-1]
+    
+    # Convert numpy array back to PIL Image
+    final_image = Image.fromarray(plotted_image_rgb)
+
+    # Encode image to Base64 to send to frontend
+    buffered = io.BytesIO()
+    final_image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
     return {
-        "damage_detected": len(bboxes) > 0,
-        "damage_type": class_names[idx.item()],
-        "confidence": round(conf.item(), 4),
-        "bboxes": bboxes
-    }
+        "damage_detected": damage_detected,
+        "damage_type": damage_type,
+        "confidence": confidence_score,
+        "annotated_image": img_str  # Base64 string of the image
+    }