Update app.py

app.py

@@ -6,27 +6,39 @@ from PIL import Image
 from ultralytics import YOLO
 
 # Load YOLOv11 Model
-model_path = "best.pt"  
+model_path = "best.pt" 
 model = YOLO(model_path)
 
 def predict(image):
     image = np.array(image)
     results = model(image)
 
+    detected_classes = set()  # Track unique detected classes
     labels = []
+
     # Draw bounding boxes and extract labels
     for result in results:
         for box in result.boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             conf = box.conf[0]
             cls = int(box.cls[0])
-            label = f"{model.names[cls]} {conf:.2f}"
-            
-            labels.append(label)  # Store detected labels
-            
+            class_name = model.names[cls]
+
+            detected_classes.add(class_name)  # Store detected class
+            label = f"{class_name} {conf:.2f}"
+            labels.append(label)
+
             cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
             cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
+    # Define possible classes (adjust based on your dataset)
+    possible_classes = {"front", "back"}
+
+    # Identify missing class if any
+    missing_classes = possible_classes - detected_classes
+    if missing_classes:
+        labels.append(f"Missing: {', '.join(missing_classes)}")
+
     return Image.fromarray(image), labels
 
 # Gradio Interface
@@ -34,7 +46,7 @@ iface = gr.Interface(
     fn=predict, 
     inputs="image", 
     outputs=["image", "text"],  # Returning both image and detected labels
-    title="YOLOv11 Object Detection"
+    title="YOLOv11 Object Detection (Front & Back Card)"
 )
 
 iface.launch()