Upload 5 files

app.py

@@ -1,49 +1,92 @@
-# Description: This is the main file to run the Gradio interface for the object detection model.
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[32]:
+
+
 from ultralytics import YOLO
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 from huggingface_hub import snapshot_download
 import os
+from torchvision import transforms
+
+classes = {0: "Defective", 1: "Good"}
 
 model_path = "best_int8_openvino_model"
 
+def load_model_local():   
+    detection_model = YOLO(model_path, task='classify') # Load the model
+    return detection_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
+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='classify')
     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
+
+    # Call the model to transform image size
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)), 
+        transforms.ToTensor(), 
+    ]) 
     
-    #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
+    source = transform(source) # Update the source image size to 224x224, 1 of 2 sizes accepted by Yolo classification model
+    
+    #result = detection_model.predict(source, conf=0.5, iou=0.6)
+    result = detection_model.predict(source) # Make prediction
+    # Get the top prediction
+    label = result[0].probs.top1 
+
+    class_names = detection_model.names  # Retrieves the class names mapping (dict-like)
+    classified_type = class_names[label]  # Map numeric label to class name
+    print (">>> Class : ", classified_type)
     
-        return out_pilimg
+    confidence = result[0].probs.top1conf # Get the top class confidence
+    print(">>> Confidence : ", confidence)
+    
+    annotated_image = pilimg.convert("RGB")
+    draw = ImageDraw.Draw(annotated_image)
+    font = ImageFont.truetype("arialbd.ttf", 30)  # Use arial.ttf for bold font
+
+    if classified_type == classes[0]:
+        draw.text((300, 10), classified_type, fill="red", font=font)
+        gr.Warning("Defect detected, BAD!.")  
+    else:
+        draw.text((300, 10), classified_type, fill="green", font=font)
+        gr.Info("No defect detected,GOOD!")  
 
-REPO_ID = "best_int8_openvino_model"  # The repo ID of the model
+    #draw.text((300, 10), classified_type, fill="red", font=font)
+        
+    return annotated_image
 
-detection_model = load_model()
+detection_model = load_model_local()
 
-title = "Classify whether the image is Defective or Good"
+title = "Detect the status of the cap, DEFECTIVE or GOOD"
 interface = gr.Interface(
                 fn=predict,
                 inputs=gr.Image(type="pil", label="Input Image"),
-                outputs=gr.Image(type="pil", label="Classified Imamge"),
+                outputs=gr.Image(type="pil", label="Classification result"),
                 title=title,
              )
 
 # Launch the interface
 interface.launch(share=True)
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+

requirements.txt

@@ -2,4 +2,5 @@ ultralytics
 gradio
 huggingface_hub
 pillow
-ffmpeg-python
+ffmpeg-python
+torchvision