Aalaa's picture
Update app.py
996e768 verified
raw
history blame
1.74 kB
import gradio as gr
import yolov7
import tempfile
import cv2
from pathlib import Path
def image_fn(image, model_path, image_size, conf_threshold, iou_threshold):
"""
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)
img = results.render()[0]
img.flags.writeable = True # Make the image writable
return img
# Use gr.components instead of gr.inputs and gr.outputs
image_interface = gr.Interface(
fn=image_fn,
inputs=[
gr.components.Image(type="pil", label="Input Image"),
gr.components.Dropdown(
choices=["Aalaa/Yolov7_Visual_Pollution_Detection"],
value="Aalaa/Yolov7_Visual_Pollution_Detection",
label="Model"
),
gr.components.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
gr.components.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
gr.components.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold")
],
outputs=gr.components.Image(type="numpy", label="Output Image"),
examples=[['image1.jpg', 'Aalaa/Yolov7_Visual_Pollution_Detection', 640, 0.25, 0.45]],
cache_examples=True,
theme='default'
)
if __name__ == "__main__":
gr.TabbedInterface(
[image_interface],
["Run on Images"],
).launch()