Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| import os | |
| model_path = "best_int8_openvino_model" | |
| TOKEN = os.getenv("authtoken") | |
| def load_model(repo_id, token): | |
| download_dir = snapshot_download(repo_id, token=token, repo_type="space") | |
| print(download_dir) | |
| path = os.path.join(download_dir, "best_int8_openvino_model") | |
| print(path) | |
| detection_model = YOLO(path, task='detect') | |
| return detection_model | |
| def predict(pilimg): | |
| source = pilimg | |
| # x = np.asarray(pilimg) | |
| # print(x.shape) | |
| result = detection_model.predict(source, conf=0.5, iou=0.6) | |
| img_bgr = result[0].plot() | |
| out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image | |
| return out_pilimg | |
| REPO_ID = "IT3103-2024S2/222012D" | |
| detection_model = load_model(REPO_ID, TOKEN) | |
| gr.Interface(fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Image(type="pil") | |
| ).launch(share=True) | |