Blessing's picture
Update app.py
c342d72 verified
raw
history blame
6.24 kB
import gradio as gr
import cv2
import requests
import os
import random
from ultralytics import YOLO
file_urls = [
'https://www.dropbox.com/scl/fi/5pavu4vvkprrtkwktvei7/DSC02373.JPG?rlkey=fpj636qtkf3vrqfxy45n2d9ii&dl=1',
'https://www.dropbox.com/scl/fi/56pbn4r3ohk85rchcvwdj/DSC02813.JPG?rlkey=jnbsidqtthk6p4ysld6o6kc4t&dl=1',
'https://www.dropbox.com/scl/fi/av9g5zbmrrzg9064zivat/image_2.jpg?rlkey=ldocvzz5lq98zffqf1lmhbhv1&dl=1',
'https://www.dropbox.com/scl/fi/izo2eqqnqzcsaxis1qrbx/IMG_7612.JPG?rlkey=6wfjaux44khtlx454ex0ng0hp&dl=1',
'https://www.dropbox.com/scl/fi/e6vgy1et6vjr61uypk5yu/VID-20230809-WA0021.mp4?rlkey=khv8rw074vezzlg8ob38bpmbx&dl=1'
]
def download_file(url, save_name):
url = url
if not os.path.exists(save_name):
file = requests.get(url)
open(save_name, 'wb').write(file.content)
for i, url in enumerate(file_urls):
if 'mp4' in file_urls[i]:
download_file(
file_urls[i],
f"video.mp4"
)
else:
download_file(
file_urls[i],
f"image_{i}.jpg"
)
model = YOLO('best.pt')
path = [['image_0.jpg'], ['image_1.jpg'], ['image_2.jpg'], ['image_3.jpg']]
# path = [['IMG_7612.JPG'], ['IMG_7678.JPG'], ['all_33.jpg'], ['all_80.jpg'],
# ['DSC02813.JPG'], ['DSC02373.JPG']]
# path = [['sc_1_0 (1) (1).JPG'], ['sc_1_0 (16) (1).JPG'],
# ['sc_1_0 (18) (1).JPG'], ['sc_1_0 (18).JPG']]
video_path = [['video.mp4']]
classes = ['alligator_cracking', 'longitudinal_cracking', 'potholes', 'ravelling']
def show_preds_image(image_path):
image = cv2.imread(image_path)
outputs = model.predict(source=image_path, agnostic_nms=True, conf=0.25, iou=0.4, imgsz=640)
results = outputs[0].cpu().numpy()
re_boxes = results.boxes.data.tolist()
class_colors = {1 : (95, 255, 54), 2: (242, 210, 100), 3: (96, 7, 70), 4:(221, 59, 41)}
random.seed(42)
# class_colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for _ in range(4)]
for i, det in enumerate(results.boxes.xyxy):
x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
class_label = int(re_boxes[i][-1])
rectangle_color = class_colors.get(class_label)
# rectangle_color = class_colors[class_label]
text_color = rectangle_color
cv2.rectangle(
image,
(int(det[0]), int(det[1])),
(int(det[2]), int(det[3])),
color=rectangle_color,
thickness=3,
lineType=cv2.LINE_AA
)
text_position = (x1, y1+100)
conf = re_boxes[i][-2]
class_name = classes[class_label]
# class_label = class_name.split('_')[0] + '\n' + class_name.split('_')[1] if '_' in class_name else class_name
cv2.putText(image, classes[class_label] + f' = {round(conf, 2)}',
text_position, cv2.FONT_HERSHEY_SIMPLEX, 1.5, text_color, 3)
# print(class_ids)
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
inputs_image = [
gr.components.Image(type="filepath", label="Input Image"),
]
outputs_image = [
gr.components.Image(type="numpy", label="Output Image"),
]
interface_image = gr.Interface(
fn=show_preds_image,
inputs=inputs_image,
outputs=outputs_image,
title="Pavement Distress Detector for developing countries",
examples=path,
cache_examples=False,
description= 'This is a demo app that takes in images or videos of Asphalt pavement surfaces and \
\n detects the following pavement distresses: \
\n \
\n Alligator cracking \
\n Longitudinal cracking \
\n Potholes \
\n Ravelling \
\n \
\n This is specifically for Inference and educational purpose.\
\n \
\n The model might ocassionaly give false outputs'
)
def show_preds_video(video_path):
cap = cv2.VideoCapture(video_path)
while(cap.isOpened()):
ret, frame = cap.read()
if ret:
frame_copy = frame.copy()
outputs = model.predict(source=frame, agnostic_nms=True, conf=0.25, iou=0.4, imgsz=640)
results = outputs[0].cpu().numpy()
re_boxes = results.boxes.data.tolist()
class_colors = {1 : (95, 255, 54), 2: (242, 210, 100), 3: (96, 7, 70), 4:(221, 59, 41)}
random.seed(42)
# class_colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for _ in range(4)]
for i, det in enumerate(results.boxes.xyxy):
x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
class_label = int(re_boxes[i][-1])
rectangle_color = class_colors.get(class_label)
# rectangle_color = class_colors[class_label]
text_color = rectangle_color
cv2.rectangle(
frame_copy,
(int(det[0]), int(det[1])),
(int(det[2]), int(det[3])),
color=rectangle_color,
thickness=2,
lineType=cv2.LINE_AA
)
text_position = (x1, y1+100)
conf = re_boxes[i][-2]
class_name = classes[class_label]
# class_label = class_name.split('_')[0] + '\n' + class_name.split('_')[1] if '_' in class_name else class_name
cv2.putText(frame_copy, classes[class_label] + f' = {round(conf, 2)}',
text_position, cv2.FONT_HERSHEY_SIMPLEX, 1.5, text_color, 3)
yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
inputs_video = [
gr.components.Video(type="filepath", label="Input Video"),
]
outputs_video = [
gr.components.Image(type="numpy", label="Output Video"),
]
interface_video = gr.Interface(
fn=show_preds_video,
inputs=inputs_video,
outputs=outputs_video,
title="Asphalt Road Pavement Distresses Detector",
examples=video_path,
cache_examples=False,
# live=True
)
gr.TabbedInterface(
[interface_image, interface_video],
tab_names=['Image inference', 'Video inference'],
).queue().launch()