AI_color / app.py
stphtan94117's picture
Update app.py
52984a4
raw
history blame
1.42 kB
import gradio as gr
import cv2
import requests
import os
from ultralytics import YOLO
# ------
file_urls = [
'https://i.imgur.com/wgZzKIk.jpg',
'https://i.imgur.com/TvIo0Nq.jpg'
]
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('crack.pt')
path = [['image_0.jpg'], ['image_1.jpg']]
video_path = [['video.mp4']]
# ------
def show_preds_image(source):
global model
res = model(source, conf=.5, iou=.5)
res_plotted = res[0].plot()
# converting BGR to RGB
result = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
return result
# ------
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="Roof Crack detector app",
examples=path,
cache_examples=False,
)
# ------
gr.TabbedInterface(
[interface_image],
tab_names=['Image inference']
).queue().launch()