Update to use ONNX

app.py

@@ -1,241 +1,115 @@
-import gradio as gr
-import tensorflow as tf
+import os
 import numpy as np
+import onnx
+import onnxruntime
 from PIL import Image, ImageDraw, ImageFont
-import os
-
-# Suppress TensorFlow logging for cleaner logs
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+import gradio as gr
 
-# Disable GPU usage explicitly to prevent TensorFlow from attempting to access GPU libraries
-os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+# Constants
+PROB_THRESHOLD = 0.5  # Minimum probability to show results
+MODEL_PATH = os.path.join("onnx", "model.onnx")
+LABELS_PATH = os.path.join("onnx", "labels.txt")
 
 # Load labels
-labels_path = os.path.join('tensorflow', 'labels.txt')
-if not os.path.exists(labels_path):
-    raise FileNotFoundError(f"Labels file not found at {labels_path}")
-with open(labels_path, 'r') as f:
-    labels = f.read().splitlines()
-
-# Function to load the frozen TensorFlow graph
-def load_frozen_graph(pb_file_path):
-    if not os.path.exists(pb_file_path):
-        raise FileNotFoundError(f"Model file not found at {pb_file_path}")
-    with tf.io.gfile.GFile(pb_file_path, 'rb') as f:
-        graph_def = tf.compat.v1.GraphDef()
-        graph_def.ParseFromString(f.read())
-    
-    with tf.Graph().as_default() as graph:
-        tf.import_graph_def(graph_def, name='')
-    return graph
-
-# Load the TensorFlow model
-MODEL_DIR = 'tensorflow'
-MODEL_PATH = os.path.join(MODEL_DIR, 'model.pb')
-graph = load_frozen_graph(MODEL_PATH)
-sess = tf.compat.v1.Session(graph=graph)
-
-# Define tensor names based on your model's outputs
-try:
-    input_tensor = graph.get_tensor_by_name('image_tensor:0')
-    detected_boxes = graph.get_tensor_by_name('detected_boxes:0')
-    detected_classes = graph.get_tensor_by_name('detected_classes:0')
-    detected_scores = graph.get_tensor_by_name('detected_scores:0')
-except KeyError as e:
-    raise KeyError(f"Tensor not found in the graph: {e}")
-
-# Define the target size based on your model's expected input
-TARGET_WIDTH = 320
-TARGET_HEIGHT = 320
-
-def preprocess_image(image):
-    """
-    Preprocess the input image:
-    - Resize to target dimensions
-    - Convert to numpy array
-    - Normalize pixel values
-    - Convert RGB to BGR if required by the model
-    """
-    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[-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
-
-def draw_boxes(image, boxes, classes, scores, threshold=0.5):
-    """
-    Draw bounding boxes and labels on the image.
-    Args:
-        image (PIL.Image): The original image.
-        boxes (np.array): Array of bounding boxes.
-        classes (np.array): Array of class IDs.
-        scores (np.array): Array of confidence scores.
-        threshold (float): Confidence threshold to filter detections.
-    Returns:
-        PIL.Image: Annotated image.
-    """
+with open(LABELS_PATH, "r") as f:
+    LABELS = f.read().strip().split("\n")
+
+class Model:
+    def __init__(self, model_filepath):
+        self.session = onnxruntime.InferenceSession(model_filepath)
+        assert len(self.session.get_inputs()) == 1
+        self.input_shape = self.session.get_inputs()[0].shape[2:]  # (H, W)
+        self.input_name = self.session.get_inputs()[0].name
+        self.input_type = {'tensor(float)': np.float32, 'tensor(float16)': np.float16}.get(
+            self.session.get_inputs()[0].type, np.float32
+        )
+        self.output_names = [o.name for o in self.session.get_outputs()]
+
+        self.is_bgr = False
+        self.is_range255 = False
+        onnx_model = onnx.load(model_filepath)
+        for metadata in onnx_model.metadata_props:
+            if metadata.key == 'Image.BitmapPixelFormat' and metadata.value == 'Bgr8':
+                self.is_bgr = True
+            elif metadata.key == 'Image.NominalPixelRange' and metadata.value == 'NominalRange_0_255':
+                self.is_range255 = True
+
+    def predict(self, image: Image.Image):
+        # Preprocess image
+        image_resized = image.resize(self.input_shape)
+        input_array = np.array(image_resized, dtype=np.float32)[np.newaxis, :, :, :]
+        input_array = input_array.transpose((0, 3, 1, 2))  # (N, C, H, W)
+        if self.is_bgr:
+            input_array = input_array[:, (2, 1, 0), :, :]
+        if not self.is_range255:
+            input_array = input_array / 255.0  # Normalize to [0,1]
+
+        # Run inference
+        outputs = self.session.run(self.output_names, {self.input_name: input_array.astype(self.input_type)})
+        return {name: outputs[i] for i, name in enumerate(self.output_names)}
+
+def draw_boxes(image: Image.Image, outputs: dict):
     draw = ImageDraw.Draw(image)
     try:
-        font = ImageFont.truetype("arial.ttf", 15)
+        font = ImageFont.truetype("arial.ttf", size=16)
     except IOError:
         font = ImageFont.load_default()
 
-    # 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
+    boxes = outputs.get('detected_boxes', [])
+    classes = outputs.get('detected_classes', [])
+    scores = outputs.get('detected_scores', [])
 
-    # Otherwise, proceed to draw bounding boxes
     for box, cls, score in zip(boxes[0], classes[0], scores[0]):
-        if score < threshold:
+        if score < PROB_THRESHOLD:
             continue
-
-        # Convert box coordinates from normalized to absolute
+        label = LABELS[int(cls)]
+        # Assuming box format: [ymin, xmin, ymax, xmax] normalized [0,1]
         ymin, xmin, ymax, xmax = box
-        left = xmin * image.width
-        right = xmax * image.width
-        top = ymin * image.height
-        bottom = ymax * image.height
-        
-        # Draw rectangle
-        draw.rectangle([(left, top), (right, bottom)], outline="red", width=2)
-        
-        # Prepare label
-        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_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 text
-        draw.text((left + 2, top - text_height - 2),
-                  label_str, fill="white", font=font)
-    
+        width, height = image.size
+        left = xmin * width
+        right = xmax * width
+        top = ymin * height
+        bottom = ymax * height
+        draw.rectangle([left, top, right, bottom], outline="red", width=2)
+        text = f"{label}: {score:.2f}"
+        text_size = draw.textsize(text, font=font)
+        draw.rectangle([left, top - text_size[1], left + text_size[0], top], fill="red")
+        draw.text((left, top - text_size[1]), text, fill="white", font=font)
     return image
 
-def predict(image):
-    """
-    Perform inference on the input image and return the annotated image.
-    Args:
-        image (PIL.Image): Uploaded image.
-    Returns:
-        PIL.Image: Annotated image with bounding boxes and labels.
-    """
-    try:
-        # Preprocess the image
-        input_array = preprocess_image(image)
+# Initialize model
+model = Model(MODEL_PATH)
 
-        # Debug prints
-        print(f"[DEBUG] Input shape to model: {input_array.shape}")
+def detect_objects(image):
+    outputs = model.predict(image)
+    annotated_image = draw_boxes(image.copy(), outputs)
 
-        # Run inference
-        boxes, classes, scores = sess.run(
-            [detected_boxes, detected_classes, detected_scores],
-            feed_dict={input_tensor: input_array}
-        )
+    # Prepare detection summary
+    detections = []
+    boxes = outputs.get('detected_boxes', [])
+    classes = outputs.get('detected_classes', [])
+    scores = outputs.get('detected_scores', [])
 
-        # 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("[DEBUG] Annotation completed.")
-        return annotated_image
-    
-    except Exception as e:
-        # Log the exception for debugging
-        print(f"Exception during prediction: {e}")
-        
-        # Return an error image with the error message
-        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
+    for box, cls, score in zip(boxes[0], classes[0], scores[0]):
+        if score < PROB_THRESHOLD:
+            continue
+        label = LABELS[int(cls)]
+        detections.append(f"{label}: {score:.2f}")
 
-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"
-description = "Upload an image of a baseball card, and JunkWaxHero will identify the set it belongs to with high accuracy."
-
-# Verify that example images exist to prevent FileNotFoundError
-example_images = ["examples/card1.jpg", "examples/card2.jpg", "examples/card3.jpg"]
-valid_examples = [img for img in example_images if os.path.exists(img)]
-if not valid_examples:
-    valid_examples = None  # Remove examples if none exist
+    detection_summary = "\n".join(detections) if detections else "No objects detected."
+
+    return annotated_image, detection_summary
 
+# Gradio Interface
 iface = gr.Interface(
-    fn=predict,
+    fn=detect_objects,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title=title,
-    description=description,
-    examples=valid_examples,
-    flagging_mode="never"  # Use new Gradio parameter
+    outputs=[gr.Image(type="pil", label="Detected Objects"), gr.Textbox(label="Detections")],
+    title="Object Detection with ONNX Model",
+    description="Upload an image to detect objects using the ONNX model.",
+    examples=["examples/card1.jpg", "examples/card2.jpg", "examples/card3.jpg"]
 )
 
 if __name__ == "__main__":
     iface.launch()
+