Update

app.py

@@ -6,65 +6,49 @@ from ultralytics import YOLO
 # Global variables to control the process
 process = False
 model = None
-cap = None
 
 def load_yolo_model():
     return YOLO('yolov8n-seg.pt')
 
-def process_video():
-    global process, model, cap
-    while process and 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):
-    global model, cap
-    model = load_yolo_model()
-    cap = cv.VideoCapture(uploaded_video.name)
-    process_video()
-
-def segment_webcam():
-    global model, cap
-    model = load_yolo_model()
-    cap = cv.VideoCapture(0)
-    process_video()
-
-def process_inputs(start_button, stop_button, mode_selection, uploaded_video):
-    global process
+def resize_frame(frame, width=1280):
+    height, width = frame.shape[:2]
+    new_height = int(height * width / float(width))
+    return cv.resize(frame, (width, new_height))
+
+def process_frame(frame):
+    global model
+    frame = resize_frame(frame)
+    start = time.perf_counter()
+    results = model(frame)
+    end = time.perf_counter()
+    segments = results[0].plot()
+    return segments, f'FPS: {int(1 // (end - start))}'
+
+def process_inputs(start_button, stop_button, mode_selection, frame, uploaded_video):
+    global process, model
     if start_button and mode_selection:
         process = True
-        if mode_selection == "Video" and uploaded_video is not None:
-            segment_video(uploaded_video)
-        elif mode_selection == "Webcam":
-            segment_webcam()
+        model = load_yolo_model()
+        if mode_selection == "Webcam" and frame is not None:
+            return process_frame(frame)
+        elif mode_selection == "Video" and uploaded_video is not None:
+            return process_frame(uploaded_video)
     elif stop_button:
         process = False
 
-start_button = gr.components.Button(label="Start")
-stop_button = gr.components.Button(label="Stop")
-mode_selection = gr.components.Radio(["Video", "Webcam"], label="Mode Selection")
-uploaded_video = gr.components.File(label="Upload Video")
+start_button = gr.inputs.Button(label="Start")
+stop_button = gr.inputs.Button(label="Stop")
+mode_selection = gr.inputs.Radio(["Webcam", "Video"], label="Mode Selection")
+frame = gr.inputs.Video(shape=(720, 1280), label="Webcam Feed")
+uploaded_video = gr.inputs.Video(shape=(720, 1280), label="Upload Video")
 
 iface = gr.Interface(
     fn=process_inputs,
-    inputs=[start_button, stop_button, mode_selection, uploaded_video],
-    outputs=None,
+    inputs=[start_button, stop_button, mode_selection, frame, uploaded_video],
+    outputs=[gr.outputs.Image(), gr.outputs.Textbox()],
     live=True,
     title="YOLO Image Segmentation",
-    description="This application uses the YOLO model to perform image segmentation on a video or webcam feed. Select a mode, upload a video (if applicable), and click 'Start' to begin. Click 'Stop' to end the process.",
+    description="This application uses the YOLO model to perform image segmentation on a webcam feed or an uploaded video. Select 'Webcam' or 'Video', upload a video (if applicable), and click 'Start' to begin. Click 'Stop' to end the process.",
     theme="huggingface"
 )