File size: 772 Bytes
911b00a
 
 
 
bf4057d
911b00a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os 

def load_model(repo_id):
    download_dir = snapshot_download(repo_id)
    print(download_dir)
    path  = os.path.join(download_dir, "best.pt")
    print(path)
    detection_model = YOLO(path, task='detect')
    return detection_model


def predict(pilimg):

    source = pilimg
    result = detection_model.predict(source, conf=0.5, iou=0.6)
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])

    return out_pilimg


REPO_ID = "Joshua-is-tired/paperclip-stickynote"
detection_model = load_model(REPO_ID)

gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Image(type="pil")).launch(share=True)