Update app.py

app.py

@@ -1,49 +1,49 @@
-# 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
-from huggingface_hub import snapshot_download
-import os
-
-model_path = "best_int8_openvino_model"
-
-
-# Load the model    
-def load_model(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 to the model
-    print(path)
-    detection_model = YOLO(path, task='classify') # 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.5) # confidence threshold, intersection over union threshold
-    
-    #print("Result: ", result)
-    if not result or len(result[0].boxes) == 0: # if no object detected
-        gr.Warning("Image not recognized for classfication!")
-    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 = "best_int8_openvino_model"  # The repo ID of the model
-
-detection_model = load_model(model_path)
-
-title = "Classify whether the image is Defective or Good"
-interface = gr.Interface(
-                fn=predict,
-                inputs=gr.Image(type="pil", label="Input Image"),
-                outputs=gr.Image(type="pil", label="Classified Imamge"),
-                title=title,
-             )
-
-# Launch the interface
-interface.launch(share=True)
+# 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
+from huggingface_hub import snapshot_download
+import os
+
+model_path = "best_int8_openvino_model"
+
+
+# Load the model    
+def load_model():
+    #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 to the model
+    #print(model_path)
+    detection_model = YOLO(model_path, task='classify') # 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.5) # confidence threshold, intersection over union threshold
+    
+    #print("Result: ", result)
+    if not result or len(result[0].boxes) == 0: # if no object detected
+        gr.Warning("Image not recognized for classfication!")
+    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 = "best_int8_openvino_model"  # The repo ID of the model
+
+detection_model = load_model()
+
+title = "Classify whether the image is Defective or Good"
+interface = gr.Interface(
+                fn=predict,
+                inputs=gr.Image(type="pil", label="Input Image"),
+                outputs=gr.Image(type="pil", label="Classified Imamge"),
+                title=title,
+             )
+
+# Launch the interface
+interface.launch(share=True)