Updated app.py

app.py

@@ -1,5 +1,3 @@
-# app.py
-
 import gradio as gr
 import tensorflow as tf
 import numpy as np
@@ -58,12 +56,12 @@ def preprocess_image(image):
     - Normalize pixel values
     - Convert RGB to BGR if required by the model
     """
-    image = image.convert('RGB')  # Ensure image is in RGB
     image = image.resize((TARGET_WIDTH, TARGET_HEIGHT))
     image_np = np.array(image).astype(np.float32)
     image_np = image_np / 255.0  # Normalize to [0,1]
-    if image_np.shape[2] == 3:
-        image_np = image_np[:, :, (2, 1, 0)]  # Convert RGB to BGR if required
+    if image_np.shape[-1] == 3:
+        # Convert RGB to BGR if required by your model
+        image_np = image_np[..., (2, 1, 0)]  
     image_np = np.expand_dims(image_np, axis=0)  # Add batch dimension
     return image_np
 
@@ -84,13 +82,17 @@ def draw_boxes(image, boxes, classes, scores, threshold=0.5):
         font = ImageFont.truetype("arial.ttf", 15)
     except IOError:
         font = ImageFont.load_default()
-    
-    detections = False  # Flag to check if any detections are above threshold
 
+    # If there are no detections at all
+    if boxes.shape[0] == 0 or boxes.shape[1] == 0:
+        # Return the original image without annotation
+        return image
+
+    # Otherwise, proceed to draw bounding boxes
     for box, cls, score in zip(boxes[0], classes[0], scores[0]):
         if score < threshold:
             continue
-        detections = True
+
         # Convert box coordinates from normalized to absolute
         ymin, xmin, ymax, xmax = box
         left = xmin * image.width
@@ -102,42 +104,24 @@ def draw_boxes(image, boxes, classes, scores, threshold=0.5):
         draw.rectangle([(left, top), (right, bottom)], outline="red", width=2)
         
         # Prepare label
-        label = f"{labels[int(cls) - 1]}: {score:.2f}"
+        cls_index = int(cls) - 1  # If your classes are 1-indexed
+        if cls_index < 0 or cls_index >= len(labels):
+            label_str = f"cls_{int(cls)}: {score:.2f}"
+        else:
+            label_str = f"{labels[cls_index]}: {score:.2f}"
         
         # Calculate text size using textbbox
-        text_bbox = draw.textbbox((0, 0), label, font=font)
+        text_bbox = draw.textbbox((0, 0), label_str, font=font)
         text_width = text_bbox[2] - text_bbox[0]
         text_height = text_bbox[3] - text_bbox[1]
         
         # Draw label background
-        draw.rectangle([(left, top - text_height - 4), (left + text_width + 4, top)], fill="red")
+        draw.rectangle([(left, top - text_height - 4),
+                        (left + text_width + 4, top)], fill="red")
         
         # Draw text
-        draw.text((left + 2, top - text_height - 2), label, fill="white", font=font)
-    
-    if not detections:
-        # Optionally, you can add text indicating no detections were found
-        try:
-            font_large = ImageFont.truetype("arial.ttf", 20)
-        except IOError:
-            font_large = ImageFont.load_default()
-        no_detect_text = "No detections found."
-        text_bbox = draw.textbbox((0, 0), no_detect_text, font=font_large)
-        text_width = text_bbox[2] - text_bbox[0]
-        text_height = text_bbox[3] - text_bbox[1]
-        draw.rectangle(
-            [
-                ((image.width - text_width) / 2 - 10, (image.height - text_height) / 2 - 10),
-                ((image.width + text_width) / 2 + 10, (image.height + text_height) / 2 + 10)
-            ],
-            fill="black"
-        )
-        draw.text(
-            ((image.width - text_width) / 2, (image.height - text_height) / 2),
-            no_detect_text,
-            fill="white",
-            font=font_large
-        )
+        draw.text((left + 2, top - text_height - 2),
+                  label_str, fill="white", font=font)
     
     return image
 
@@ -150,26 +134,30 @@ def predict(image):
         PIL.Image: Annotated image with bounding boxes and labels.
     """
     try:
-        print("Starting prediction...")
         # Preprocess the image
         input_array = preprocess_image(image)
-        print(f"Preprocessed image shape: {input_array.shape}")
-        
+
+        # Debug prints
+        print(f"[DEBUG] Input shape to model: {input_array.shape}")
+
         # Run inference
         boxes, classes, scores = sess.run(
             [detected_boxes, detected_classes, detected_scores],
             feed_dict={input_tensor: input_array}
         )
-        print(f"Inference completed. Number of detections: {len(boxes[0])}")
-        
-        # Check if detections are present
-        if boxes.shape[1] == 0:
-            print("No detections returned by the model.")
-        
+
+        # Debug prints
+        print(f"[DEBUG] boxes shape: {boxes.shape}, classes shape: {classes.shape}, scores shape: {scores.shape}")
+
+        # Check if the output arrays have zero detections
+        # e.g. boxes could have shape (1, 0, 4) if no detections are found
+        if boxes.size == 0 or classes.size == 0 or scores.size == 0:
+            print("[DEBUG] No detections returned by the model.")
+            return _draw_no_detection_message(image)
+
         # Annotate the image with bounding boxes and labels
         annotated_image = draw_boxes(image.copy(), boxes, classes, scores, threshold=0.5)
-        print("Annotated image created successfully.")
-        
+        print("[DEBUG] Annotation completed.")
         return annotated_image
     
     except Exception as e:
@@ -177,34 +165,57 @@ def predict(image):
         print(f"Exception during prediction: {e}")
         
         # Return an error image with the error message
-        error_image = Image.new('RGB', (500, 500), color=(255, 0, 0))
-        draw = ImageDraw.Draw(error_image)
-        try:
-            font = ImageFont.truetype("arial.ttf", 20)
-        except IOError:
-            font = ImageFont.load_default()
-        error_text = "Error during prediction."
-        
-        # Calculate text size using textbbox
-        text_bbox = draw.textbbox((0, 0), error_text, font=font)
-        text_width = text_bbox[2] - text_bbox[0]
-        text_height = text_bbox[3] - text_bbox[1]
-        
-        # Center the text
-        draw.rectangle(
-            [
-                ((500 - text_width) / 2 - 10, (500 - text_height) / 2 - 10),
-                ((500 + text_width) / 2 + 10, (500 + text_height) / 2 + 10)
-            ],
-            fill="black"
-        )
-        draw.text(
-            ((500 - text_width) / 2, (500 - text_height) / 2),
-            error_text,
-            fill="white",
-            font=font
-        )
-        return error_image
+        return _draw_error_message()
+
+def _draw_no_detection_message(image):
+    """Draws a simple 'No detections found' message on the image."""
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)
+    except IOError:
+        font = ImageFont.load_default()
+    message = "No detections found."
+    text_bbox = draw.textbbox((0, 0), message, font=font)
+    text_width = text_bbox[2] - text_bbox[0]
+    text_height = text_bbox[3] - text_bbox[1]
+
+    # Center the message
+    x = (image.width - text_width) / 2
+    y = (image.height - text_height) / 2
+    draw.rectangle(
+        [(x - 10, y - 10), (x + text_width + 10, y + text_height + 10)],
+        fill="black"
+    )
+    draw.text((x, y), message, fill="white", font=font)
+    return image
+
+def _draw_error_message():
+    """Creates a red image with a centered error message."""
+    error_image = Image.new('RGB', (500, 500), color=(255, 0, 0))
+    draw = ImageDraw.Draw(error_image)
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)
+    except IOError:
+        font = ImageFont.load_default()
+    error_text = "Error during prediction."
+    text_bbox = draw.textbbox((0, 0), error_text, font=font)
+    text_width = text_bbox[2] - text_bbox[0]
+    text_height = text_bbox[3] - text_bbox[1]
+    
+    draw.rectangle(
+        [
+            ((500 - text_width) / 2 - 10, (500 - text_height) / 2 - 10),
+            ((500 + text_width) / 2 + 10, (500 + text_height) / 2 + 10)
+        ],
+        fill="black"
+    )
+    draw.text(
+        ((500 - text_width) / 2, (500 - text_height) / 2),
+        error_text,
+        fill="white",
+        font=font
+    )
+    return error_image
 
 # Define Gradio interface using the new API
 title = "JunkWaxHero 🦸‍♂️ - Baseball Card Set Identifier"
@@ -223,7 +234,7 @@ iface = gr.Interface(
     title=title,
     description=description,
     examples=valid_examples,
-    flagging_mode="never"  # Updated parameter
+    flagging_mode="never"  # Use new Gradio parameter
 )
 
 if __name__ == "__main__":