utils.py

import os
import requests
from io import BytesIO
import numpy as np
from PIL import Image
import yolov5
from yolov5.utils.plots import Annotator, colors
import gradio as gr


def load_model(model_path, img_size=640):
    HF_TOKEN = os.getenv("HF_TOKEN")
    if HF_TOKEN is not None:  # assume SECRET variable is set
        model = yolov5.load(model_path, hf_token=HF_TOKEN)
    else:
        model = yolov5.load(model_path)
    model.img_size = img_size  # add img_size attribute
    return model


def load_image_from_url(url):
    if not url:  # empty or None
        return gr.Image(interactive=True)
    try:
        response = requests.get(url, timeout=5)
        image = Image.open(BytesIO(response.content))
    except Exception as e:
        raise gr.Error("Unable to load image from URL") from e
    return image.convert("RGB")


def inference(model, image):
    results = model(image, size=model.img_size)
    annotator = Annotator(np.asarray(image))
    for *box, _, cls in reversed(results.pred[0]):
        # label = f'{model.names[int(cls)]} {conf:.2f}'
        # print(f'{cls} {conf:.2f} {box}')
        annotator.box_label(box, "", color=colors(cls, True))
    return annotator.im