Create app.py

app.py

@@ -0,0 +1,414 @@
+import argparse
+import colorsys
+import os
+
+import gradio as gr
+import numpy as np
+import onnxruntime as ort
+import pandas as pd
+from PIL import Image, ImageDraw
+
+
+def resize_with_aspect_ratio(image, size, interpolation=Image.BILINEAR):
+    """Resizes an image while maintaining aspect ratio and pads it."""
+    original_width, original_height = image.size
+    ratio = min(size / original_width, size / original_height)
+    new_width = int(original_width * ratio)
+    new_height = int(original_height * ratio)
+    image = image.resize((new_width, new_height), interpolation)
+
+    # Create a new image with the desired size and paste the resized image onto it
+    new_image = Image.new("RGB", (size, size))
+    new_image.paste(image, ((size - new_width) // 2, (size - new_height) // 2))
+    return new_image, ratio, (size - new_width) // 2, (size - new_height) // 2
+
+
+def generate_colors(num_classes):
+    """Generate a list of distinct colors for different classes."""
+    # Generate evenly spaced hues
+    hsv_tuples = [(x / num_classes, 0.8, 0.9) for x in range(num_classes)]
+
+    # Convert to RGB
+    colors = []
+    for hsv in hsv_tuples:
+        rgb = colorsys.hsv_to_rgb(*hsv)
+        # Convert to 0-255 range and to tuple
+        colors.append(tuple(int(255 * x) for x in rgb))
+
+    return colors
+
+
+def draw(images, labels, boxes, scores, ratios, paddings, thrh=0.4, class_names=None):
+    result_images = []
+
+    # Generate colors for classes
+    num_classes = (
+        len(class_names) if class_names else 91
+    )  # Use length of class_names if available, otherwise default to COCO's 91 classes
+    colors = generate_colors(num_classes)
+
+    for i, im in enumerate(images):
+        draw = ImageDraw.Draw(im)
+        scr = scores[i]
+        lab = labels[i][scr > thrh]
+        box = boxes[i][scr > thrh]
+        scr = scr[scr > thrh]
+
+        ratio = ratios[i]
+        pad_w, pad_h = paddings[i]
+
+        for lbl, bb in zip(lab, box):
+            # Get color for this class
+            class_idx = int(lbl)
+            color = colors[class_idx % len(colors)]
+
+            # Convert RGB to hex for PIL
+            hex_color = "#{:02x}{:02x}{:02x}".format(*color)
+
+            # Adjust bounding boxes according to the resizing and padding
+            bb = [
+                (bb[0] - pad_w) / ratio,
+                (bb[1] - pad_h) / ratio,
+                (bb[2] - pad_w) / ratio,
+                (bb[3] - pad_h) / ratio,
+            ]
+
+            # Draw rectangle with class-specific color
+            draw.rectangle(bb, outline=hex_color, width=3)
+
+            # Use class name if available, otherwise use class index
+            if class_names and class_idx < len(class_names):
+                label_text = f"{class_names[class_idx]} {scr[lab == lbl][0]:.2f}"
+            else:
+                label_text = f"Class {class_idx} {scr[lab == lbl][0]:.2f}"
+
+            # Draw text background
+            text_size = draw.textbbox((0, 0), label_text, font=None)
+            text_width = text_size[2] - text_size[0]
+            text_height = text_size[3] - text_size[1]
+
+            # Draw text background rectangle
+            draw.rectangle(
+                [bb[0], bb[1] - text_height - 4, bb[0] + text_width + 4, bb[1]],
+                fill=hex_color,
+            )
+
+            # Draw text in white or black depending on color brightness
+            brightness = (color[0] * 299 + color[1] * 587 + color[2] * 114) / 1000
+            text_color = "black" if brightness > 128 else "white"
+
+            # Draw text
+            draw.text(
+                (bb[0] + 2, bb[1] - text_height - 2), text=label_text, fill=text_color
+            )
+
+        result_images.append(im)
+    return result_images
+
+
+def load_model(model_path):
+    """
+    Load an ONNX model for inference.
+
+    Args:
+        model_path: Path to the ONNX model file
+
+    Returns:
+        tuple: (session, error_message)
+    """
+    providers = ["CPUExecutionProvider"]
+
+    try:
+        sess = ort.InferenceSession(model_path, providers=providers)
+        print(f"Using device: {ort.get_device()}")
+        return sess, None
+    except Exception as e:
+        return None, f"Error creating inference session: {e}"
+
+
+def load_class_names(class_names_path):
+    """
+    Load class names from a text file.
+
+    Args:
+        class_names_path: Path to a text file with class names (one per line)
+
+    Returns:
+        list: Class names or None if loading failed
+    """
+    if not class_names_path or not os.path.exists(class_names_path):
+        return None
+
+    try:
+        with open(class_names_path, "r") as f:
+            class_names = [line.strip() for line in f.readlines()]
+        print(f"Loaded {len(class_names)} class names")
+        return class_names
+    except Exception as e:
+        print(f"Error loading class names: {e}")
+        return None
+
+
+def prepare_image(image):
+    """
+    Prepare image for inference by converting to PIL and resizing.
+
+    Args:
+        image: Input image (PIL or numpy array)
+
+    Returns:
+        tuple: (resized_image, original_image, ratio, padding)
+    """
+    # Convert to PIL image if needed
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image).convert("RGB")
+
+    # Resize image while preserving aspect ratio
+    resized_image, ratio, pad_w, pad_h = resize_with_aspect_ratio(image, 640)
+
+    return resized_image, image, ratio, (pad_w, pad_h)
+
+
+def run_inference(session, image):
+    """
+    Run inference on the prepared image.
+
+    Args:
+        session: ONNX runtime session
+        image: Prepared PIL image
+
+    Returns:
+        tuple: (labels, boxes, scores)
+    """
+    # Get original image dimensions
+    orig_height, orig_width = image.size[1], image.size[0]
+    # Convert to int64 as expected by the model
+    orig_size = np.array([[orig_height, orig_width]], dtype=np.int64)
+
+    # Convert PIL image to numpy array and normalize to 0-1 range
+    im_data = np.array(image, dtype=np.float32) / 255.0
+    # Transpose from HWC to CHW format
+    im_data = im_data.transpose(2, 0, 1)
+    # Add batch dimension
+    im_data = np.expand_dims(im_data, axis=0)
+
+    output = session.run(
+        output_names=None,
+        input_feed={"images": im_data, "orig_target_sizes": orig_size},
+    )
+
+    return output  # labels, boxes, scores
+
+
+def count_objects(labels, scores, confidence_threshold, class_names):
+    """
+    Count detected objects by class.
+
+    Args:
+        labels: Detection labels
+        scores: Detection confidence scores
+        confidence_threshold: Minimum confidence threshold
+        class_names: List of class names
+
+    Returns:
+        dict: Counts of objects by class
+    """
+    object_counts = {}
+    for i, score_batch in enumerate(scores):
+        for j, score in enumerate(score_batch):
+            if score >= confidence_threshold:
+                label = labels[i][j]
+                class_name = (
+                    class_names[label]
+                    if class_names and label < len(class_names)
+                    else f"Class {label}"
+                )
+                object_counts[class_name] = object_counts.get(class_name, 0) + 1
+
+    return object_counts
+
+
+def create_status_message(object_counts):
+    """
+    Create a status message with object counts.
+
+    Args:
+        object_counts: Dictionary of object counts by class
+
+    Returns:
+        str: Formatted status message
+    """
+    status_message = "Detection completed successfully\n\nObjects detected:"
+    if object_counts:
+        for class_name, count in object_counts.items():
+            status_message += f"\n- {class_name}: {count}"
+    else:
+        status_message += "\n- No objects detected above confidence threshold"
+
+    return status_message
+
+
+def create_bar_data(object_counts):
+    """
+    Create data for the bar plot visualization.
+
+    Args:
+        object_counts: Dictionary of object counts by class
+
+    Returns:
+        DataFrame: Data for bar plot
+    """
+    if object_counts:
+        # Sort by count in descending order
+        sorted_counts = sorted(object_counts.items(), key=lambda x: x[1], reverse=True)
+        class_names_list = [item[0] for item in sorted_counts]
+        counts_list = [item[1] for item in sorted_counts]
+        # Create a pandas DataFrame for the bar plot
+        return pd.DataFrame({"Class": class_names_list, "Count": counts_list})
+    else:
+        return pd.DataFrame({"Class": ["No objects detected"], "Count": [0]})
+
+
+def predict(image, model_path, class_names_path, confidence_threshold):
+    """
+    Main prediction function that orchestrates the detection pipeline.
+
+    Args:
+        image: Input image
+        model_path: Path to ONNX model
+        class_names_path: Path to class names file
+        confidence_threshold: Detection confidence threshold
+
+    Returns:
+        tuple: (result_image, status_message, bar_data)
+    """
+    # Load model
+    session, error = load_model(model_path)
+    if error:
+        return None, error, None
+
+    # Load class names
+    class_names = load_class_names(class_names_path)
+
+    try:
+        # Prepare image
+        resized_image, original_image, ratio, padding = prepare_image(image)
+
+        # Run inference
+        labels, boxes, scores = run_inference(session, resized_image)
+
+        # Draw detections on the original image
+        result_images = draw(
+            [original_image],
+            labels,
+            boxes,
+            scores,
+            [ratio],
+            [padding],
+            thrh=confidence_threshold,
+            class_names=class_names,
+        )
+
+        # Count objects by class
+        object_counts = count_objects(labels, scores, confidence_threshold, class_names)
+
+        # Create status message
+        status_message = create_status_message(object_counts)
+
+        # Create bar plot data
+        bar_data = create_bar_data(object_counts)
+
+        return result_images[0], status_message, bar_data
+    except Exception as e:
+        return None, f"Error during inference: {e}", None
+
+
+def build_interface(model_path, class_names_path):
+    """
+    Build the Gradio interface components.
+
+    Args:
+        model_path: Path to the ONNX model
+        class_names_path: Path to the class names file
+
+    Returns:
+        gr.Blocks: The Gradio demo interface
+    """
+    with gr.Blocks(title="Blood Cell Detection") as demo:
+        gr.Markdown("# Blood Cell Detection")
+        gr.Markdown("Upload an image to detect blood cells. The model can detect 3 types of blood cells: red blood cells, white blood cells and platelets.")
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Input Image")
+                confidence = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.4,
+                    step=0.05,
+                    label="Confidence Threshold",
+                )
+                submit_btn = gr.Button("Detect Objects")
+
+            with gr.Column():
+                output_image = gr.Image(type="pil", label="Detection Result")
+
+                with gr.Row():
+                    output_message = gr.Textbox(label="Status")
+
+                    count_plot = gr.BarPlot(
+                        y="Class",
+                        x="Count",
+                        title="Object Counts",
+                        tooltip=["Class", "Count"],
+                        height=300,
+                        orientation="h",
+                        label_title="Object Counts",
+                    )
+
+        # Set up the click event inside the Blocks context
+        submit_btn.click(
+            fn=predict,
+            inputs=[
+                input_image,
+                gr.State(model_path),
+                gr.State(class_names_path),
+                confidence,
+            ],
+            outputs=[output_image, output_message, count_plot],
+        )
+
+        return demo
+
+
+def launch_demo(args):
+    """
+    Launch the Gradio demo with the specified arguments.
+
+    Args:
+        args: Command-line arguments
+    """
+    demo = build_interface(args.onnx, args.class_names)
+
+    # Launch the demo
+    demo.launch(share=args.share)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Gradio demo for object detection with ONNX Runtime"
+    )
+    parser.add_argument(
+        "--onnx", type=str, required=True, help="Path to the ONNX model file"
+    )
+    parser.add_argument(
+        "--class-names",
+        type=str,
+        help="Path to a text file with class names (one per line)",
+    )
+    parser.add_argument(
+        "--share", action="store_true", help="Create a shareable link for the demo"
+    )
+    args = parser.parse_args()
+
+    launch_demo(args)