app.py

import gradio as gr
import yolov7
import tempfile
import cv2
import numpy as np
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)
    
    # Convert PIL image to NumPy array if necessary
    if isinstance(results.render()[0], np.ndarray):
        img = results.render()[0]
    else:
        img = np.array(results.render()[0])
    
    # Ensure the image array is writable
    img.flags.writeable = True
    
    return img

# Use gr instead of gr.inputs and gr.outputs
image_interface = gr.Interface(
    fn=image_fn,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Dropdown(
            choices=["Aalaa/Yolov7_Visual_Pollution_Detection"],
            value="Aalaa/Yolov7_Visual_Pollution_Detection",
            label="Model"
        ),
        gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
        gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
        gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold")
    ],
    outputs=gr.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()