asif00 commited on
Commit
5abf1f3
·
1 Parent(s): ebaceb4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -9
app.py CHANGED
@@ -3,13 +3,17 @@ import gradio as gr
3
  import cv2 as cv
4
  from ultralytics import YOLO
5
 
 
 
 
 
6
 
7
  def load_yolo_model():
8
  return YOLO('yolov8n-seg.pt')
9
 
10
-
11
- def process_video(model, cap):
12
- while cap.isOpened():
13
  ret, image = cap.read()
14
  if not ret:
15
  break
@@ -26,26 +30,28 @@ def process_video(model, cap):
26
  cap.release()
27
  cv.destroyAllWindows()
28
 
29
-
30
  def segment_video(uploaded_video):
 
31
  model = load_yolo_model()
32
  cap = cv.VideoCapture(uploaded_video.name)
33
- process_video(model, cap)
34
-
35
 
36
  def segment_webcam():
 
37
  model = load_yolo_model()
38
  cap = cv.VideoCapture(0)
39
- process_video(model, cap)
40
-
41
 
42
  def process_inputs(start_button, stop_button, mode_selection, uploaded_video):
 
43
  if start_button:
 
44
  if mode_selection == "Video" and uploaded_video:
45
  segment_video(uploaded_video)
46
  elif mode_selection == "Webcam":
47
  segment_webcam()
48
-
 
49
 
50
  iface = gr.Interface(
51
  fn=process_inputs,
 
3
  import cv2 as cv
4
  from ultralytics import YOLO
5
 
6
+ # Global variables to control the process
7
+ process = False
8
+ model = None
9
+ cap = None
10
 
11
  def load_yolo_model():
12
  return YOLO('yolov8n-seg.pt')
13
 
14
+ def process_video():
15
+ global process, model, cap
16
+ while process and cap.isOpened():
17
  ret, image = cap.read()
18
  if not ret:
19
  break
 
30
  cap.release()
31
  cv.destroyAllWindows()
32
 
 
33
  def segment_video(uploaded_video):
34
+ global model, cap
35
  model = load_yolo_model()
36
  cap = cv.VideoCapture(uploaded_video.name)
37
+ process_video()
 
38
 
39
  def segment_webcam():
40
+ global model, cap
41
  model = load_yolo_model()
42
  cap = cv.VideoCapture(0)
43
+ process_video()
 
44
 
45
  def process_inputs(start_button, stop_button, mode_selection, uploaded_video):
46
+ global process
47
  if start_button:
48
+ process = True
49
  if mode_selection == "Video" and uploaded_video:
50
  segment_video(uploaded_video)
51
  elif mode_selection == "Webcam":
52
  segment_webcam()
53
+ elif stop_button:
54
+ process = False
55
 
56
  iface = gr.Interface(
57
  fn=process_inputs,