himisir commited on
Commit
f1913e7
·
1 Parent(s): 91aa865

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -39
app.py CHANGED
@@ -5,46 +5,49 @@ from ultralytics import YOLO
5
 
6
 
7
  def load_yolo_model():
8
- return YOLO()
9
 
10
 
11
  def process_video(model, cap):
12
-     while cap.isOpened():
13
-         ret, image = cap.read()
14
-         if not ret:
15
-             break
16
-         start = time.perf_counter()
17
-         results = model(image)
18
-         end = time.perf_counter()
19
-         segments = results[0].plot()
20
-         cv.putText(segments, f'FPS: {int(1 // (end - start))}', (10, 30),
21
-                    cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
22
-         cv.imshow('Image Segmentation', segments)
23
-         key = cv.waitKey(1)
24
-         if key & 0xFF == ord('q'):
25
-             break
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 wrapper_for_segment_video(uploaded_video):
43
- segment_video(uploaded_video)
 
 
 
44
 
45
 
46
  iface = gr.Interface(
47
- fn=None,
48
  inputs=[gr.inputs.Checkbox(label="Segment Video"),
49
  gr.inputs.Checkbox(label="Segment Webcam"),
50
  gr.inputs.File(label="Upload Video")],
@@ -52,17 +55,4 @@ iface = gr.Interface(
52
  live=True
53
  )
54
 
55
-
56
- def update_fn(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
57
- if segment_video_checkbox and uploaded_video:
58
- iface.fn = wrapper_for_segment_video
59
- iface.args = [uploaded_video]
60
- elif segment_webcam_checkbox:
61
- iface.fn = segment_webcam
62
- iface.args = []
63
- else:
64
- iface.fn = None
65
-
66
-
67
- iface.update(update_fn)
68
  iface.launch()
 
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
16
+ start = time.perf_counter()
17
+ results = model(image)
18
+ end = time.perf_counter()
19
+ segments = results[0].plot()
20
+ cv.putText(segments, f'FPS: {int(1 // (end - start))}', (10, 30),
21
+ cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
22
+ cv.imshow('Image Segmentation', segments)
23
+ key = cv.waitKey(1)
24
+ if key & 0xFF == ord('q'):
25
+ break
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(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
43
+ if segment_video_checkbox and uploaded_video:
44
+ segment_video(uploaded_video)
45
+ elif segment_webcam_checkbox:
46
+ segment_webcam()
47
 
48
 
49
  iface = gr.Interface(
50
+ fn=process_inputs,
51
  inputs=[gr.inputs.Checkbox(label="Segment Video"),
52
  gr.inputs.Checkbox(label="Segment Webcam"),
53
  gr.inputs.File(label="Upload Video")],
 
55
  live=True
56
  )
57
 
 
 
 
 
 
 
 
 
 
 
 
 
 
58
  iface.launch()