asif00 commited on
Commit
ce7f5f5
·
1 Parent(s): 0508e39

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -38
app.py CHANGED
@@ -5,63 +5,62 @@ 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")],
51
- outputs=None,
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)
 
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
+     inputs=[gr.inputs.Checkbox(label="Segment Video"),
48
+             gr.inputs.Checkbox(label="Segment Webcam"),
49
+             gr.inputs.File(label="Upload Video")],
50
+     outputs=None,
51
+     live=True
 
52
  )
53
 
54
 
55
  def update_fn(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
56
+     if segment_video_checkbox and uploaded_video:
57
+         iface.fn = wrapper_for_segment_video
58
+         iface.args = [uploaded_video]
59
+     elif segment_webcam_checkbox:
60
+         iface.fn = segment_webcam
61
+         iface.args = []
62
+     else:
63
+         iface.fn = None
64
 
65
 
66
  iface.update(update_fn)