Merge branch 'main' of https://huggingface.co/spaces/asif00/Realtime_Image_Segmentation_Using_YOLO_and_OpenCV

app.py

@@ -5,49 +5,46 @@ from ultralytics import YOLO
 
 
 def load_yolo_model():
-    return YOLO('yolov8n-seg.pt')
+    return YOLO()
 
 
 def process_video(model, cap):
-    while cap.isOpened():
-        ret, image = cap.read()
-        if not ret:
-            break
-        start = time.perf_counter()
-        results = model(image)
-        end = time.perf_counter()
-        segments = results[0].plot()
-        cv.putText(segments, f'FPS: {int(1 // (end - start))}', (10, 30),
-                   cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-        cv.imshow('Image Segmentation', segments)
-        key = cv.waitKey(1)
-        if key & 0xFF == ord('q'):
-            break
-    cap.release()
-    cv.destroyAllWindows()
+    while cap.isOpened():
+        ret, image = cap.read()
+        if not ret:
+            break
+        start = time.perf_counter()
+        results = model(image)
+        end = time.perf_counter()
+        segments = results[0].plot()
+        cv.putText(segments, f'FPS: {int(1 // (end - start))}', (10, 30),
+                   cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+        cv.imshow('Image Segmentation', segments)
+        key = cv.waitKey(1)
+        if key & 0xFF == ord('q'):
+            break
+    cap.release()
+    cv.destroyAllWindows()
 
 
 def segment_video(uploaded_video):
-    model = load_yolo_model()
-    cap = cv.VideoCapture(uploaded_video.name)
-    process_video(model, cap)
+    model = load_yolo_model()
+    cap = cv.VideoCapture(uploaded_video.name)
+    process_video(model, cap)
 
 
 def segment_webcam():
-    model = load_yolo_model()
-    cap = cv.VideoCapture(0)
-    process_video(model, cap)
+    model = load_yolo_model()
+    cap = cv.VideoCapture(0)
+    process_video(model, cap)
 
 
-def process_inputs(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
-    if segment_video_checkbox and uploaded_video:
-        segment_video(uploaded_video)
-    elif segment_webcam_checkbox:
-        segment_webcam()
+def wrapper_for_segment_video(uploaded_video):
+    segment_video(uploaded_video)
 
 
 iface = gr.Interface(
-    fn=process_inputs,
+    fn=None,
     inputs=[gr.inputs.Checkbox(label="Segment Video"),
             gr.inputs.Checkbox(label="Segment Webcam"),
             gr.inputs.File(label="Upload Video")],
@@ -55,4 +52,17 @@ iface = gr.Interface(
     live=True
 )
 
+
+def update_fn(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
+    if segment_video_checkbox and uploaded_video:
+        iface.fn = wrapper_for_segment_video
+        iface.args = [uploaded_video]
+    elif segment_webcam_checkbox:
+        iface.fn = segment_webcam
+        iface.args = []
+    else:
+        iface.fn = None
+
+
+iface.update(update_fn)
 iface.launch()