Aalaa's picture
Update app.py
76433a2 verified
raw
history blame contribute delete
No virus
1.83 kB
import gradio as gr
import yolov7
import subprocess
import tempfile
import time
from pathlib import Path
import uuid
import cv2
import gradio as gr
def image_fn(
image: gr.inputs.Image = None,
model_path: gr.inputs.Dropdown = None,
image_size: gr.inputs.Slider = 640,
conf_threshold: gr.inputs.Slider = 0.25,
iou_threshold: gr.inputs.Slider = 0.45,
):
"""
YOLOv7 inference function
Args:
image: Input image
model_path: Path to the model
image_size: Image size
conf_threshold: Confidence threshold
iou_threshold: IOU threshold
Returns:
Rendered image
"""
model = yolov7.load(model_path, device="cpu", hf_model=True, trace=False)
model.conf = conf_threshold
model.iou = iou_threshold
results = model([image], size=image_size)
return results.render()[0]
image_interface = gr.Interface(
fn=image_fn,
inputs=[
gr.inputs.Image(type="pil", label="Input Image"),
gr.inputs.Dropdown(
choices=[
"Aalaa/Yolov7_Visual_Pollution_Detection",
],
default="Aalaa/Yolov7_Visual_Pollution_Detection",
label="Model",
)
#gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size")
#gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
#gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold")
],
outputs=gr.outputs.Image(type="filepath", label="Output Image"),
examples=[['image1.jpg', 'Aalaa/Yolov7_Visual_Pollution_Detection', 640, 0.25, 0.45]],
cache_examples=True,
theme='huggingface',
)
if __name__ == "__main__":
gr.TabbedInterface(
[image_interface],
["Run on Images"],
).launch()