Spaces:
Sleeping
Sleeping
File size: 1,541 Bytes
a8ab365 9572155 a8ab365 |
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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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 the snapshot
download_dir = snapshot_download(repo_id)
print(f"Downloaded snapshot directory: {download_dir}")
# Check the contents of the directory
for root, dirs, files in os.walk(download_dir):
print(f"Found files: {files}")
# Check for the correct model file and path
# Ensure the path points to the correct location where the model is saved
model_path = os.path.join(download_dir, "best_int8_openvino_model")
# Check if the model file exists at the expected location
if not os.path.exists(model_path):
print(f"Model file not found at {model_path}")
return None
# Load the model using YOLO
detection_model = YOLO(model_path, task='detect')
return detection_model
def predict(pilimg):
result = detection_model.predict(pilimg, conf=0.5, iou=0.6)
img_bgr = result[0].plot() # Get image with predictions
out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # Convert BGR to RGB
return out_pilimg
# Set the repo ID
REPO_ID = "Lesterchia174/Monkey_Durian"
# Load the model
detection_model = load_model(REPO_ID)
if detection_model:
gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil")
).launch(share=True)
else:
print("Model loading failed. Check the model path or file structure.")
|