dwkurnie commited on
Commit
89fbcb1
·
1 Parent(s): 7261123
Files changed (1) hide show
  1. app.py +60 -24
app.py CHANGED
@@ -1,11 +1,38 @@
1
  import gradio as gr
2
  import cv2
 
 
 
3
  from ultralytics import YOLO
4
 
5
- # Load YOLO model
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  model = YOLO('best.pt')
 
 
7
 
8
- # Function to perform object detection on an image
9
  def show_preds_image(image_path):
10
  image = cv2.imread(image_path)
11
  outputs = model.predict(source=image_path)
@@ -21,9 +48,23 @@ def show_preds_image(image_path):
21
  )
22
  return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
23
 
24
- # Function to perform object detection on a video stream
25
- def show_preds_video(video_stream):
26
- cap = cv2.VideoCapture(video_stream.name)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
  while(cap.isOpened()):
28
  ret, frame = cap.read()
29
  if ret:
@@ -41,28 +82,23 @@ def show_preds_video(video_stream):
41
  )
42
  yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
43
 
44
- # Define interfaces for image and video inference
45
- interface_image = gr.Interface(
46
- fn=show_preds_image,
47
- inputs=gr.inputs.Image(type="file", label="Upload Image"),
48
- outputs=gr.outputs.Image(type="numpy", label="Output Image"),
49
- title="Pothole Detector (Image Input)",
50
- )
51
 
 
 
 
 
52
  interface_video = gr.Interface(
53
  fn=show_preds_video,
54
- inputs=gr.inputs.Video(type="webcam", label="Webcam Input"),
55
- outputs=gr.outputs.Image(type="numpy", label="Output Image"),
56
- title="Pothole Detector (Webcam Input)",
 
 
57
  )
58
 
59
- # Launch tabbed interface for both image and video inference
60
- gr.Interface(
61
  [interface_image, interface_video],
62
- title="Pothole Detector",
63
- description="Detect potholes using YOLOv8 on images and webcam streams.",
64
- examples=[
65
- ["path/to/image.jpg"],
66
- ["path/to/video.mp4"]
67
- ]
68
- ).launch()
 
1
  import gradio as gr
2
  import cv2
3
+ import requests
4
+ import os
5
+
6
  from ultralytics import YOLO
7
 
8
+ file_urls = [
9
+ 'https://www.dropbox.com/s/b5g97xo901zb3ds/pothole_example.jpg?dl=1',
10
+ 'https://www.dropbox.com/s/86uxlxxlm1iaexa/pothole_screenshot.png?dl=1',
11
+ 'https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?dl=1'
12
+ ]
13
+
14
+ def download_file(url, save_name):
15
+ url = url
16
+ if not os.path.exists(save_name):
17
+ file = requests.get(url)
18
+ open(save_name, 'wb').write(file.content)
19
+
20
+ for i, url in enumerate(file_urls):
21
+ if 'mp4' in file_urls[i]:
22
+ download_file(
23
+ file_urls[i],
24
+ f"video.mp4"
25
+ )
26
+ else:
27
+ download_file(
28
+ file_urls[i],
29
+ f"image_{i}.jpg"
30
+ )
31
+
32
  model = YOLO('best.pt')
33
+ path = [['image_0.jpg'], ['image_1.jpg']]
34
+ video_path = [['video.mp4']]
35
 
 
36
  def show_preds_image(image_path):
37
  image = cv2.imread(image_path)
38
  outputs = model.predict(source=image_path)
 
48
  )
49
  return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
50
 
51
+ inputs_image = [
52
+ gr.components.Image(type="filepath", label="Input Image"),
53
+ ]
54
+ outputs_image = [
55
+ gr.components.Image(type="numpy", label="Output Image"),
56
+ ]
57
+ interface_image = gr.Interface(
58
+ fn=show_preds_image,
59
+ inputs=inputs_image,
60
+ outputs=outputs_image,
61
+ title="Pothole detector app",
62
+ examples=path,
63
+ cache_examples=False,
64
+ )
65
+
66
+ def show_preds_video(video_path):
67
+ cap = cv2.VideoCapture(video_path)
68
  while(cap.isOpened()):
69
  ret, frame = cap.read()
70
  if ret:
 
82
  )
83
  yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
84
 
85
+ inputs_video = [
86
+ gr.components.Video(type="filepath", label="Input Video"),
 
 
 
 
 
87
 
88
+ ]
89
+ outputs_video = [
90
+ gr.components.Image(type="numpy", label="Output Image"),
91
+ ]
92
  interface_video = gr.Interface(
93
  fn=show_preds_video,
94
+ inputs=inputs_video,
95
+ outputs=outputs_video,
96
+ title="Pothole detector",
97
+ examples=video_path,
98
+ cache_examples=False,
99
  )
100
 
101
+ gr.TabbedInterface(
 
102
  [interface_image, interface_video],
103
+ tab_names=['Image inference', 'Video inference']
104
+ ).queue().launch()