Upload 3 files

app.py

@@ -1,9 +1,4 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-# In[4]:
-
-
+# Description: This is the main file to run the Gradio interface for the object detection model.
 from ultralytics import YOLO
 from PIL import Image
 import gradio as gr
@@ -12,28 +7,35 @@ import os
 
 model_path = "best_int8_openvino_model"
 
+# Example paths for Gradio
+image_examples = [["DurianMangosteen1.jpg"], ["DurianMangosteen2.jpg"]]
+
+# Load the model    
 def load_model(repo_id):
-    download_dir = snapshot_download(repo_id)
+    download_dir = snapshot_download(repo_id) # download the model from the Hugging Face Hub
     print(download_dir)
-    path  = os.path.join(download_dir, "best_int8_openvino_model")
+    path  = os.path.join(download_dir, "best_int8_openvino_model") # path to the model
     print(path)
-    detection_model = YOLO(path, task='detect')
+    detection_model = YOLO(path, task='detect') # load the model
     return detection_model
 
-
+# Predict the image
 def predict(pilimg):
-
     source = pilimg
     # x = np.asarray(pilimg)
     # print(x.shape)
-    result = detection_model.predict(source, conf=0.4, iou=0.6)
-    img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
+    result = detection_model.predict(source, conf=0.4, iou=0.6) # confidence threshold, intersection over union threshold
     
-    return out_pilimg
-
+    #print("Result: ", result)
+    if not result or len(result[0].boxes) == 0: # if no object detected
+        gr.Warning("No object detected in the image!")
+    else:
+        img_bgr = result[0].plot() # plot the image
+        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
+    
+        return out_pilimg
 
-REPO_ID = "ITI107-2024S2/8035531F"
+REPO_ID = "ITI107-2024S2/8035531F"  # The repo ID of the model
 detection_model = load_model(REPO_ID)
 
 title = "Detect Durian and Mangosteen (King and Queen of Fruits) In The Image"
@@ -42,14 +44,13 @@ interface = gr.Interface(
                 inputs=gr.Image(type="pil", label="Input Image"),
                 outputs=gr.Image(type="pil", label="Object Detected Image"),
                 title=title,
+                examples=image_examples,
              )
 
 # Launch the interface
 interface.launch(share=True)
 
 
-# In[ ]:
-