| 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) | |
| model_path = os.path.join(download_dir, "best.pt") | |
| model = YOLO(model_path) | |
| return model | |
| def predict(pil_image): | |
| result = detection_model.predict(pil_image, conf=0.7, iou=0.6) | |
| img_bgr = result[0].plot() | |
| output = Image.fromarray(img_bgr[..., ::-1]) | |
| return output | |
| REPO_ID = "cedrickjohn/your_model_repo" | |
| detection_model = load_model(REPO_ID) | |
| gr.Interface(fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Image(type="pil"), | |
| title="Battery Key Detection").launch() | |