File size: 993 Bytes
876f0dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
from detect import *
from pathlib import Path
# Get the current script's directory
script_dir = Path(__file__).resolve().parent

# Define the relative path to the weights file
weights_path = script_dir / 'runs/train/exp17/weights/best.pt'


def yolo_inference(image):
    import tempfile
    temp_dir = tempfile.TemporaryDirectory()
    temp_path = Path(temp_dir.name) / "input_image.jpg"
    output_path = Path(temp_dir.name) / "output_image.jpg"
    image.save(temp_path)

    run(weights=ROOT / weights_path, source=str(temp_path), project=str(temp_dir.name), name='result', exist_ok=True)
    
    result_image_path = list((Path(temp_dir.name) / 'result').glob('*.jpg'))[0]
    result_image = Image.open(result_image_path)
    
    return result_image

interface = gr.Interface(
    fn=yolo_inference,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="YOLO Object Detection",
    description="Upload an image to detect objects using YOLO."
)

interface.launch()