Update services/under_construction/culvert_check.py

services/under_construction/culvert_check.py

@@ -1,18 +1,19 @@
 from ultralytics import YOLO
 import cv2
-import os
+import gradio as gr
 
-MODEL_PATH = "yolov8m.pt"  # Replace with culvert-specific model if available
+MODEL_PATH = "yolov8m.pt"
 model = YOLO(MODEL_PATH)
 
 def run_culvert_check(video_file):
     if video_file is None:
-        return "Please upload a video."
-    
+        return "Please upload a video.", None
+
     cap = cv2.VideoCapture(video_file)
     results_summary = []
-
+    detected_frame = None
     frame_idx = 0
+
     while True:
         success, frame = cap.read()
         if not success:
@@ -24,10 +25,15 @@ def run_culvert_check(video_file):
         results = model.predict(source=frame, save=False, conf=0.3)
         if results:
             for r in results:
-                if hasattr(r, 'boxes') and r.boxes is not None:
+                if hasattr(r, 'boxes') and r.boxes is not None and len(r.boxes) > 0:
                     count = len(r.boxes)
                     results_summary.append(f"Frame {frame_idx}: {count} culvert features detected.")
+                    annotated = r.plot()
+                    detected_frame = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
+                    cap.release()
+                    return "\n".join(results_summary), detected_frame
+
         frame_idx += 1
-    cap.release()
 
-    return "\n".join(results_summary) if results_summary else "No culvert features found."
+    cap.release()
+    return "No culvert features found.", None