Upload app.py

app.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[8]:
+
+
+import tensorflow as tf
+from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
+from tensorflow.keras.preprocessing import image
+from ultralytics import YOLO
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import gradio as gr
+from huggingface_hub import snapshot_download
+import os
+from torchvision import transforms
+
+# Define the class labels
+classes = {0: "Defective", 1: "Good"}
+
+model_path = "ResNet50_Classification.h5" # Trained RestNet50 model
+
+best_yolo_model = "best.pt" # Trained YOLOv8 detection model
+
+classification_model = tf.keras.models.load_model('ResNet50_Classification.h5') 
+
+detection_model = YOLO(best_yolo_model, task='detect')
+
+
+# Define the image preprocessing function
+def preprocess_imageXX(image):
+    image = image.resize((224, 224))  # Resize to the input size of ResNet50
+    image = np.array(image)  # Convert to numpy array
+    image = preprocess_input(image)  # Preprocess for ResNet50
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    return image
+    
+def preprocess_image(pilimg):
+    img = pilimg.resize((224, 224))  # Resize to the input size of ResNet50    
+    img_array = image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension    
+    return img_array
+
+def classify_image(pilimg):
+    img_array = preprocess_image(pilimg)  # Preprocess the input image
+    classify_result = classification_model.predict(img_array)[0][0]  # Get prediction probability 
+    print(">>> Result : ", classify_result)
+
+    predicted_class = "Good" if classify_result >= 0.5 else "Defective"
+    print(">>> predicted_class : ", predicted_class)
+
+    return predicted_class
+
+def detect_defect(img):
+    detection_result = detection_model.predict(img, conf=0.4, iou=0.5)
+     
+    return detection_result
+
+   
+def process_image(pilimg):
+    # Perform classification first, then perform detection if Defective
+    classification = classify_image(pilimg)
+
+    if classification == "Good":
+        out_pilimg = pilimg.convert("RGB")
+        draw = ImageDraw.Draw(out_pilimg)
+        font = ImageFont.truetype("arialbd.ttf", 30)  # Use arial.ttf for bold font
+
+        draw.text((250, 10), "Good", fill="green", font=font)
+        gr.Info("No defect detected,GOOD!", duration=3)  
+
+    else:  # Defective
+        detection_result = detect_defect(pilimg)
+        img_bgr = detection_result[0].plot()
+        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
+
+        draw = ImageDraw.Draw(out_pilimg)
+        font = ImageFont.truetype("arialbd.ttf", 30)  # Use arial.ttf for bold font
+
+        draw.text((250, 10), "Defective", fill="red", font=font)
+        
+        detections = detection_result[0].boxes.data  # Get detections
+        if len(detections) > 0:
+            gr.Warning("Defect detected, BAD!", duration=3)
+        else:
+            gr.Warning("Classified as Defective but defect cannot be detected, ERROR!")
+            
+    return out_pilimg
+
+title = "Detect the status of the cap, DEFECTIVE or GOOD"
+interface = gr.Interface(
+                fn=process_image,
+                inputs=gr.Image(type="pil", label="Input Image"),
+                outputs=gr.Image(type="pil", label="Classification result"),
+                title=title,
+             )
+
+# Launch the interface
+interface.launch(share=True)
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+
+
+# In[37]:
+
+
+from ultralytics import YOLO
+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
+
+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
+
+def predict(pilimg):
+    source = pilimg
+
+    # Call the model to transform image size
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)), 
+        transforms.ToTensor(), 
+    ]) 
+    
+    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.55, 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)
+    
+    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!")  
+
+    #draw.text((300, 10), classified_type, fill="red", font=font)
+        
+    return annotated_image
+
+detection_model = load_model_local()
+
+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="Classification result"),
+                title=title,
+             )
+
+# Launch the interface
+interface.launch(share=True)
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+