Update app.py

app.py

@@ -6,12 +6,13 @@ from PIL import Image, ImageEnhance
 from ultralytics import YOLO
 import json
 
-
 model_path = "best.pt" 
 model = YOLO(model_path)
 
 def preprocess_image(image):
-    image = Image.fromarray(np.array(image))
+    """Preprocesses the image: enhances sharpness, contrast, brightness, and resizes it."""
+    if isinstance(image, np.ndarray):  # Ensure it's a PIL image
+        image = Image.fromarray(image)
 
     image = ImageEnhance.Sharpness(image).enhance(2.0)  # Increase sharpness
     image = ImageEnhance.Contrast(image).enhance(1.5)   # Increase contrast
@@ -25,12 +26,8 @@ def preprocess_image(image):
 
     return image
 
-def imageRotation(image):
-    """Dummy function for now."""
-    return image
-
 def vision_ai_api(image, label):
-    """Dummy function simulating API call. Returns dummy JSON response."""
+    """Dummy function simulating an API call. Returns dummy JSON response."""
     return {
         "label": label,
         "extracted_data": {
@@ -41,19 +38,29 @@ def vision_ai_api(image, label):
     }
 
 def predict(image):
-    # Convert PIL image to NumPy array
-    if isinstance(image, Image.Image):
-        image = np.array(image)
+    """Runs YOLO object detection on the input image and processes detected regions."""
+    # Ensure image is PIL format before preprocessing
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
     image = preprocess_image(image)  # Apply preprocessing
 
-    results = model(image, conf=0.80)
+    # Convert image to NumPy array for YOLO model
+    image_np = np.array(image)
 
+    # Run YOLO prediction
+    results = model(image_np, conf=0.80)
+    
     detected_classes = set()
     labels = []
     cropped_images = {}
 
+    # Ensure results contain boxes
     for result in results:
+        if result.boxes is None or len(result.boxes) == 0:
+            print("No objects detected.")
+            continue
+        
         for box in result.boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             conf = box.conf[0]
@@ -65,16 +72,17 @@ def predict(image):
             detected_classes.add(class_name)
             labels.append(f"{class_name} {conf:.2f}")
 
-            # Ensure bounding boxes are within the image
-            height, width = image.shape[:2]  # ✅ This now works
+            # Ensure bounding boxes are within the image dimensions
+            height, width = image_np.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]  # Crop the detected region
-            cropped_pil = Image.fromarray(cropped)  # Convert back to PIL
+            # Crop the detected region
+            cropped = image_np[y1:y2, x1:x2]
+            cropped_pil = Image.fromarray(cropped)  # Convert to PIL for API
 
             # Call API
             api_response = vision_ai_api(cropped_pil, class_name)
@@ -83,6 +91,7 @@ def predict(image):
                 "api_response": json.dumps(api_response, indent=4)
             }
 
+    # Ensure outputs exist even if no detections were made
     if not cropped_images:
         return None, "No front detected", None, "No back detected", ["No valid detections"]
 
@@ -95,11 +104,9 @@ def predict(image):
     )
 
 
-
-# Gradio Interface
 iface = gr.Interface(
     fn=predict, 
-    inputs="image", 
+    inputs=gr.Image(type="pil"),  # Ensure input is PIL image
     outputs=["image", "text"],  
     title="License Field Detection (Front & Back Card)",
     description="Detect front & back of a license card, crop the images, and call Vision AI API separately for each."