|
from ultralytics import YOLO |
|
from PIL import Image |
|
import gradio as gr |
|
from huggingface_hub import snapshot_download |
|
import os |
|
|
|
model_path = "" |
|
|
|
def load_model(repo_id): |
|
download_dir = snapshot_download(repo_id) |
|
print(download_dir) |
|
path = os.path.join(download_dir, "") |
|
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 = "ITI107-2024S2/7821983W-lowboonsiong" |
|
detection_model = load_model(REPO_ID) |
|
|
|
gr.Interface(fn=predict, |
|
inputs=gr.Image(type="pil"), |
|
outputs=gr.Image(type="pil") |
|
).launch(share=True) |
|
|
|
|