Update app.py

app.py

@@ -1,34 +1,21 @@
 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_int8_openvino_model")
-    print(path)
-    detection_model = YOLO(path, task='detect')
-    return detection_model
 
+# Load model directly from file in the Space
+detection_model = YOLO("best.pt")
 
+# Prediction function
 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
-    
+    results = detection_model.predict(pilimg, conf=0.5, iou=0.6)
+    img_bgr = results[0].plot()
+    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
     return out_pilimg
 
-
-REPO_ID = "KendrickTan/231941k"
-detection_model = load_model(REPO_ID)
-
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil")
-             ).launch(share=True)
+# Gradio interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Mask Detection Demo"
+).launch()