Refactor Gradio Demo Script for Clarity and Functionality

- Removed unused absolute path definitions for model and class names.
- Simplified the `draw` and `prepare_image` functions for better readability.
- Enhanced error handling in the `load_model` function.
- Updated the Gradio interface to streamline model path selection and improve user experience.
- Cleaned up commented code and improved variable naming for clarity.

Author: dnth

scripts/gradio_demo.py

@@ -7,17 +7,12 @@
 import numpy as np
 import onnxruntime as ort
 import pandas as pd
-from PIL import Image, ImageDraw
+from PIL import Image
 import cv2
 
 ort.preload_dlls()
 
 
-# Use absolute paths instead of relative paths
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-MODEL_PATH = os.path.join(BASE_DIR, "models/deim-blood-cell-detection_nano.onnx")
-CLASS_NAMES_PATH = os.path.join(BASE_DIR, "models/classes.txt")
-
 def generate_colors(num_classes):
     """Generate a list of distinct colors for different classes."""
     # Generate evenly spaced hues
@@ -52,18 +47,22 @@ def draw(images, labels, boxes, scores, scales, paddings, thrh=0.4, class_names=
         # Scale boxes from padded size to original image size
         scale = scales[i]
         x_offset, y_offset = paddings[i]
-        
-        valid_boxes[:, [0, 2]] = (valid_boxes[:, [0, 2]] - x_offset) / scale  # x coordinates
-        valid_boxes[:, [1, 3]] = (valid_boxes[:, [1, 3]] - y_offset) / scale  # y coordinates
+
+        valid_boxes[:, [0, 2]] = (
+            valid_boxes[:, [0, 2]] - x_offset
+        ) / scale  # x coordinates
+        valid_boxes[:, [1, 3]] = (
+            valid_boxes[:, [1, 3]] - y_offset
+        ) / scale  # y coordinates
 
         # Draw boxes
         for label, box, score in zip(valid_labels, valid_boxes, valid_scores):
             class_idx = int(label)
             color = colors[class_idx % len(colors)]
-            
+
             # Convert coordinates to integers
             box = [int(coord) for coord in box]
-            
+
             # Draw rectangle
             cv2.rectangle(im, (box[0], box[1]), (box[2], box[3]), color, 2)
 
@@ -125,35 +124,14 @@ def load_model(model_path):
         print(f"Loading model from: {model_path}")
         if not os.path.exists(model_path):
             return None, f"Model file not found at: {model_path}"
-            
+
         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 get_classes_path(custom_path, default_path):
-    """
-    Get class names file path.
-    
-    Args:
-        custom_path: Custom path to class names file
-        default_path: Default path to class names file
-        
-    Returns:
-        Path to a class names file
-    """
-    if not custom_path:
-        return default_path
-    
-    # Treat as a file path
-    if os.path.exists(custom_path):
-        return custom_path
-    
-    return default_path
-
-
 def load_class_names(class_names_path):
     """
     Load class names from a text file.
@@ -180,31 +158,23 @@ def load_class_names(class_names_path):
 def prepare_image(image, target_size=640):
     """
     Prepare image for inference by converting to PIL and resizing with padding.
-
-    Args:
-        image: Input image (PIL or numpy array)
-        target_size: Target size for resizing (default: 640)
-
-    Returns:
-        tuple: (model_input, original_image, scale, padding)
     """
     # Convert to numpy array if PIL Image
     if isinstance(image, Image.Image):
         image = np.array(image)
-    
-    # Calculate scaling and padding
+
     height, width = image.shape[:2]
     scale = target_size / max(height, width)
-    new_height = int(height * scale)
-    new_width = int(width * scale)
+    new_height, new_width = int(height * scale), int(width * scale)
 
     # Calculate padding
-    y_offset = (target_size - new_height) // 2
-    x_offset = (target_size - new_width) // 2
+    y_offset, x_offset = [(target_size - dim) // 2 for dim in (new_height, new_width)]
 
     # Create model input with padding
     model_input = np.zeros((target_size, target_size, 3), dtype=np.uint8)
-    model_input[y_offset:y_offset + new_height, x_offset:x_offset + new_width] = cv2.resize(image, (new_width, new_height))
+    model_input[y_offset : y_offset + new_height, x_offset : x_offset + new_width] = (
+        cv2.resize(image, (new_width, new_height))
+    )
 
     return model_input, image, scale, (x_offset, y_offset)
 
@@ -230,7 +200,9 @@ def run_inference(session, image, target_size=640):
         dtype=np.float32,
     )
     im_data = np.expand_dims(im_data, axis=0)  # Add batch dimension
-    orig_size = np.array([[target_size, target_size]], dtype=np.int64)  # Use padded size
+    orig_size = np.array(
+        [[target_size, target_size]], dtype=np.int64
+    )  # Use padded size
 
     # Get input name and run inference
     input_name = session.get_inputs()[0].name
@@ -320,7 +292,7 @@ def predict(image, model_path, class_names_path, confidence_threshold, image_siz
     model_load_start = time.time()
     session, error = load_model(model_path)
     model_load_time = time.time() - model_load_start
-    
+
     if error:
         return None, error, None
 
@@ -337,10 +309,10 @@ def predict(image, model_path, class_names_path, confidence_threshold, image_siz
         inference_start = time.time()
         outputs = run_inference(session, model_input, image_size)
         inference_time = time.time() - inference_start
-        
+
         if not outputs or len(outputs) < 3:
             return None, "Error: Model output is invalid", None
-            
+
         labels, boxes, scores = outputs
 
         # Draw detections
@@ -363,17 +335,18 @@ def predict(image, model_path, class_names_path, confidence_threshold, image_siz
         # Create status message with timing information
         status_message = create_status_message(object_counts)
         status_message += "\n\nLatency Information:"
-        status_message += f"\n- Model Loading: {model_load_time*1000:.1f}ms"
-        status_message += f"\n- Preprocessing: {preprocess_time*1000:.1f}ms"
-        status_message += f"\n- Inference: {inference_time*1000:.1f}ms"
-        status_message += f"\n- Postprocessing: {postprocess_time*1000:.1f}ms"
-        status_message += f"\n- Total Time: {(model_load_time + preprocess_time + inference_time + postprocess_time)*1000:.1f}ms"
-        
+        status_message += f"\n- Model Loading: {model_load_time * 1000:.1f}ms"
+        status_message += f"\n- Preprocessing: {preprocess_time * 1000:.1f}ms"
+        status_message += f"\n- Inference: {inference_time * 1000:.1f}ms"
+        status_message += f"\n- Postprocessing: {postprocess_time * 1000:.1f}ms"
+        status_message += f"\n- Total Time: {(model_load_time + preprocess_time + inference_time + postprocess_time) * 1000:.1f}ms"
+
         bar_data = create_bar_data(object_counts)
 
         return result_images[0], status_message, bar_data
     except Exception as e:
         import traceback
+
         error_details = traceback.format_exc()
         print(f"Error during inference: {error_details}")
         return None, f"Error during inference: {str(e)}", None
@@ -382,37 +355,15 @@ def predict(image, model_path, class_names_path, confidence_threshold, image_siz
 def build_interface(model_path, class_names_path, example_images=None):
     """
     Build the Gradio interface components.
-
-    Args:
-        model_path: Path to the ONNX model
-        class_names_path: Path to the class names file
-        example_images: List of example image paths
-
-    Returns:
-        gr.Blocks: The Gradio demo interface
     """
     with gr.Blocks(title="DEIMKit Detection") as demo:
         gr.Markdown("# DEIMKit Detection")
-        gr.Markdown("Configure the model and run inference on an image.")
-        
-        # Add model selection
-        with gr.Accordion("Model Settings", open=False):
-            with gr.Row():
-                custom_model_path = gr.File(
-                    label="Custom Model File (ONNX)",
-                    file_types=[".onnx"],
-                    file_count="single"
-                )
-                custom_classes_path = gr.File(
-                    label="Custom Classes File (TXT)",
-                    file_types=[".txt"],
-                    file_count="single"
-                )
+        gr.Markdown("Upload an image and run inference.")
 
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="Input Image")
-                
+
                 with gr.Row():
                     confidence = gr.Slider(
                         minimum=0.1,
@@ -421,24 +372,23 @@ def build_interface(model_path, class_names_path, example_images=None):
                         step=0.01,
                         label="Confidence Threshold",
                     )
-                    
+
                     image_size = gr.Slider(
                         minimum=32,
                         maximum=1920,
                         value=640,
                         step=32,
                         label="Image Size",
-                        info="Select image size for inference (larger = slower but potentially more accurate)"
+                        info="Select image size for inference (larger = slower but potentially more accurate)",
                     )
-                
+
                 submit_btn = gr.Button("Run Inference", variant="primary")
 
             with gr.Column():
                 output_image = gr.Image(type="pil", label="Detection Result")
 
                 with gr.Row(equal_height=True):
                     output_message = gr.Textbox(label="Status")
-
                     count_plot = gr.BarPlot(
                         x="Class",
                         y="Count",
@@ -448,71 +398,43 @@ def build_interface(model_path, class_names_path, example_images=None):
                         orientation="h",
                         label_title="Object Counts",
                     )
-        
+
         # Add examples component if example images are provided
         if example_images:
             gr.Examples(
                 examples=example_images,
                 inputs=input_image,
             )
 
-        # Function to handle model path selection
-        def get_model_path(custom_file, default_path):
-            if custom_file is not None:
-                return custom_file.name
-            return default_path
-            
-        def get_classes_path(custom_file, default_path):
-            if custom_file is not None:
-                return custom_file.name
-            return default_path
-
-        # Set up the click event inside the Blocks context
+        # Set up the click event
         submit_btn.click(
-            fn=lambda img, custom_model, custom_classes, conf, img_size: predict(
+            fn=lambda img, conf, img_size: predict(
                 img,
-                get_model_path(custom_model, model_path),
-                get_classes_path(custom_classes, class_names_path),
+                model_path,
+                class_names_path,
                 conf,
-                img_size
+                img_size,
             ),
-            inputs=[
-                input_image,
-                custom_model_path,
-                custom_classes_path,
-                confidence,
-                image_size,
-            ],
+            inputs=[input_image, confidence, image_size],
             outputs=[output_image, output_message, count_plot],
         )
 
         with gr.Row():
             with gr.Column():
-                gr.HTML("<div style='text-align: center; margin: 0 auto;'>Created by <a href='https://dicksonneoh.com' target='_blank'>Dickson Neoh</a>.</div>")
+                gr.HTML(
+                    "<div style='text-align: center; margin: 0 auto;'>Created by <a href='https://dicksonneoh.com' target='_blank'>Dickson Neoh</a>.</div>"
+                )
 
         return demo
 
 
 def parse_args():
     """Parse command line arguments."""
-    parser = argparse.ArgumentParser(description='DEIMKit Detection Demo')
-    parser.add_argument(
-        '--model', 
-        type=str,
-        default=MODEL_PATH,
-        help='Path to ONNX model file'
-    )
-    parser.add_argument(
-        '--classes', 
-        type=str,
-        default=CLASS_NAMES_PATH,
-        help='Path to class names file'
-    )
+    parser = argparse.ArgumentParser(description="DEIMKit Detection Demo")
+    parser.add_argument("--model", type=str, required=True, help="Path to ONNX model file")
+    parser.add_argument("--classes", type=str, required=True, help="Path to class names file")
     parser.add_argument(
-        '--examples',
-        type=str,
-        default=os.path.join(BASE_DIR, "examples"),
-        help='Path to directory containing example images'
+        "--examples", type=str, help="Path to directory containing example images (optional)"
     )
     return parser.parse_args()
 
@@ -522,27 +444,18 @@ def launch_demo():
     Launch the Gradio demo with model and class names paths from command line arguments.
     """
     args = parse_args()
-    
-    # Create examples directory if it doesn't exist
-    examples_dir = args.examples
-    if not os.path.exists(examples_dir):
-        os.makedirs(examples_dir)
-        print(f"Created examples directory at {examples_dir}")
-    
-    # Get list of example images
+
     example_images = []
-    if os.path.exists(examples_dir):
+    if args.examples and os.path.exists(args.examples):
         example_images = [
-            os.path.join(examples_dir, f) 
-            for f in os.listdir(examples_dir) 
-            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
+            os.path.join(args.examples, f)
+            for f in os.listdir(args.examples)
+            if f.lower().endswith((".png", ".jpg", ".jpeg"))
         ]
-        print(f"Found {len(example_images)} example images in {examples_dir}")
-    
+        print(f"Found {len(example_images)} example images in {args.examples}")
+
     demo = build_interface(args.model, args.classes, example_images)
-    
-    # Launch the demo without the examples parameter
-    demo.launch(share=False, inbrowser=True)  # Set share=True if you want to create a shareable link
+    demo.launch(share=False, inbrowser=True)
 
 
 if __name__ == "__main__":