app.py

import subprocess
import tempfile
import time
from pathlib import Path
import cv2
import gradio as gr
import torch
from inferer import Inferer

pipeline = Inferer("nateraw/yolov6s", device='cuda')
print(f"GPU on? {'🟢' if pipeline.device.type != 'cpu' else '🔴'}")


def fn_image(image, conf_thres, iou_thres):
    return pipeline(image, conf_thres, iou_thres)


gr.Interface(
    fn_image,
    inputs=[
        gr.Image(source='webcam', streaming=True),
        gr.Slider(0, 1, value=0.5, label="Confidence Threshold"),
        gr.Slider(0, 1, value=0.5, label="IOU Threshold"),
    ],
    outputs=gr.Image(type="file"),
    live=True,
    title="YOLOv8",
    allow_flagging=False,
    allow_screenshot=False,
)