Update app.py

app.py

@@ -44,9 +44,10 @@ def vision_ai_api(image, label):
     }
 
 def predict(image):
-    image = preprocess_image(image)  # Apply preprocessing
+    image = preprocess_image(image)
+    
+    results = model(image, conf=0.80)
 
-    results = model(image, conf=0.85)
     detected_classes = set()
     labels = []
     cropped_images = {}
@@ -57,37 +58,38 @@ def predict(image):
             conf = box.conf[0]
             cls = int(box.cls[0])
             class_name = model.names[cls]
+            
+            print(f"Detected: {class_name} ({conf:.2f}) at [{x1}, {y1}, {x2}, {y2}]")
 
             detected_classes.add(class_name)
             labels.append(f"{class_name} {conf:.2f}")
 
-            # Crop detected region
+            # Ensure bounding boxes are within the image
+            height, width = image.shape[:2]
+            x1, y1, x2, y2 = max(0, x1), max(0, y1), min(width, x2), min(height, y2)
+
+            if x1 >= x2 or y1 >= y2:
+                print("Invalid bounding box, skipping.")
+                continue
+
             cropped = image[y1:y2, x1:x2]
             cropped_pil = Image.fromarray(cropped)
 
-            # Call Vision AI API separately for front & back
+            # Call API
             api_response = vision_ai_api(cropped_pil, class_name)
+            cropped_images[class_name] = {"image": cropped_pil, "api_response": json.dumps(api_response, indent=4)}
 
-            # Store cropped images & API response
-            cropped_images[class_name] = {
-                "image": cropped_pil,
-                "api_response": json.dumps(api_response, indent=4)
-            }
-
-    # Identify missing classes
-    possible_classes = {"front", "back"}
-    missing_classes = possible_classes - detected_classes
-    if missing_classes:
-        labels.append(f"Missing: {', '.join(missing_classes)}")
-
-    # Prepare Gradio outputs (separate front & back images and responses)
-    front_image = cropped_images.get("front", {}).get("image", None)
-    back_image = cropped_images.get("back", {}).get("image", None)
+    if not cropped_images:
+        return None, "No front detected", None, "No back detected", ["No valid detections"]
 
-    front_response = cropped_images.get("front", {}).get("api_response", "{}")
-    back_response = cropped_images.get("back", {}).get("api_response", "{}")
+    return (
+        cropped_images.get("front", {}).get("image", None),
+        cropped_images.get("front", {}).get("api_response", "{}"),
+        cropped_images.get("back", {}).get("image", None),
+        cropped_images.get("back", {}).get("api_response", "{}"),
+        labels
+    )
 
-    return front_image, front_response, back_image, back_response, labels
 
 # Gradio Interface
 iface = gr.Interface(