Update app.py

app.py

@@ -3,24 +3,29 @@ import torch
 import cv2
 from ultralytics import YOLO
 
-# Safe load method to handle custom YOLO class during deserialization
+# Load YOLO models
 def safe_load_yolo_model(path):
-    # Add necessary safe globals to allow the detection model class during loading
     torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel'])
     return YOLO(path)
 
-# Load YOLO models
+# Load the models
 model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
 model_best = safe_load_yolo_model('./data/best.pt')
 
+# Class names for YOLO model (replace with actual class names used in your YOLO model)
+yolo_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe']
+
+# Class names for best.pt model (assumed classes for crack and pothole)
+best_classes = ['Crack', 'Pothole']
+
 def process_video(video):
-    # Read video input (Gradio gives you a file-like object)
+    # Open the video using OpenCV
     cap = cv2.VideoCapture(video)
     fps = cap.get(cv2.CAP_PROP_FPS)
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    # Create a VideoWriter object to save the output video
+    # Create VideoWriter to save output video
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4
     out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))
 
@@ -29,27 +34,30 @@ def process_video(video):
         if not ret:
             break
         
-        # Use both YOLO models for detection
+        # Detect with YOLOv11 (general object detection model)
         results_yolo11 = model_yolo11(frame)
+        # Detect with best.pt (specialized model for cracks and potholes)
         results_best = model_best(frame)
-        
-        # Combine the results from both models
-        # For simplicity, we will overlay bounding boxes and labels from both models
+
+        # Draw bounding boxes and labels for YOLOv11 (General Object Detection)
         for result in results_yolo11:
             boxes = result.boxes
             for box in boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                class_id = int(box.cls[0])  # Class index for YOLO
+                label = f"YOLO: {yolo_classes[class_id]} - {box.conf[0]:.2f}"  # Map class_id to class name
                 cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                label = f"YOLOv11: {box.cls[0]} - {box.conf[0]:.2f}"
-                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
 
+        # Draw bounding boxes and labels for best.pt (Crack and Pothole detection)
         for result in results_best:
             boxes = result.boxes
             for box in boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                class_id = int(box.cls[0])  # Class index for best.pt
+                label = f"Best: {best_classes[class_id]} - {box.conf[0]:.2f}"  # Map class_id to specific labels
                 cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
-                label = f"Best: {box.cls[0]} - {box.conf[0]:.2f}"
-                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
 
         # Write the processed frame to the output video
         out.write(frame)