Update app.py

app.py

@@ -8,63 +8,74 @@ def safe_load_yolo_model(path):
     torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel'])
     return YOLO(path)
 
-# Load the models
-model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
-model_best = safe_load_yolo_model('./data/best2.pt')
+# Dictionary of model paths
+model_paths = {
+    'YOLOv11': './data/yolo11n.pt',
+    'Crack & Pothole Detector': './data/best2.pt',
+    'Bridge Inspector': './data/bridge.pt',
+    'Road Surface Detector': './data/road.pt',
+    'Pipe Detector': './data/pipe.pt'
+}
 
-def process_video(video):
-    # Open the video using OpenCV
+# Load models into memory
+models = {name: safe_load_yolo_model(path) for name, path in model_paths.items()}
+
+# Assign colors for each model
+model_colors = {
+    'YOLOv11': (0, 255, 0),
+    'Crack & Pothole Detector': (255, 0, 0),
+    'Bridge Inspector': (0, 0, 255),
+    'Road Surface Detector': (255, 255, 0),
+    'Pipe Detector': (255, 0, 255)
+}
+
+def process_video(video, selected_model):
     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 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))
+    out = cv2.VideoWriter('output_video.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+
+    use_models = models if selected_model == 'All' else {selected_model: models[selected_model]}
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
-        # 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)
 
-        # 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
-                # Use model's built-in class names
-                label = f"{model_yolo11.names[class_id]} - {box.conf[0]:.2f}"
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+        for model_name, model in use_models.items():
+            results = model(frame)
 
-        # 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
-                # Use model's built-in class names for best.pt
-                label = f"{model_best.names[class_id]} - {box.conf[0]:.2f}"
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
-                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+            for result in results:
+                for box in result.boxes:
+                    x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                    class_id = int(box.cls[0])
+                    label = f"{model.names[class_id]} - {box.conf[0]:.2f}"
+                    color = model_colors.get(model_name, (0, 255, 255))
+                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
 
-        # Write the processed frame to the output video
         out.write(frame)
 
     cap.release()
     out.release()
-
     return 'output_video.mp4'
 
-# Gradio interface
-iface = gr.Interface(fn=process_video, inputs=gr.Video(), outputs=gr.Video(), live=True)
+# Gradio Interface
+iface = gr.Interface(
+    fn=process_video,
+    inputs=[
+        gr.Video(label="Upload a Video"),
+        gr.Dropdown(
+            choices=["All"] + list(model_paths.keys()),
+            label="Select Model(s)",
+            value="All"
+        )
+    ],
+    outputs=gr.Video(label="Processed Output"),
+    live=False,
+    title="Multi-Model YOLOv8 Video Inference"
+)
 
-# Launch the app
 iface.launch()