Create onnx_runtime_draw_bboxes.py

onnx_runtime_draw_bboxes.py

@@ -0,0 +1,77 @@
+import onnxruntime as ort
+import numpy as np
+import cv2
+import os
+
+# Load the ONNX model
+model_path = "model.onnx"
+session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+
+# Folder paths
+input_folder = "input_images"  # Folder containing images
+output_folder = "output_images"  # Folder to save processed images
+os.makedirs(output_folder, exist_ok=True)
+
+# Define model input shape
+MODEL_INPUT_SIZE = (320, 320)  # Change based on your model input size
+
+# Define class labels (update based on your model)
+CLASS_NAMES = ["person", "car", "truck", "bicycle", "dog", "cat"]  # Update accordingly
+
+def preprocess_image(image_path):
+    """ Preprocess image for ONNX model input """
+    image = cv2.imread(image_path)
+    original_image = image.copy()  # Save original for later drawing
+    image = cv2.resize(image, MODEL_INPUT_SIZE)  # Resize to model input size
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert to RGB
+    image = image.astype(np.float32) / 255.0  # Normalize
+    image = np.transpose(image, (2, 0, 1))  # Convert to CHW format
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    return image, original_image
+
+def postprocess_output(output, orig_image):
+    """ Post-process ONNX model output """
+    height, width, _ = orig_image.shape
+    
+    # Extract boxes, scores, and class indices
+    boxes = output[0]  # Adjust key names if needed
+    scores = output[1]
+    class_indices = output[2]
+
+    # Draw bounding boxes
+    for i in range(len(scores)):
+        if scores[i] > 0.5:  # Confidence threshold
+            x1, y1, x2, y2 = boxes[i]  # Get box coordinates
+            x1, y1, x2, y2 = int(x1 * width), int(y1 * height), int(x2 * width), int(y2 * height)  # Scale box
+            
+            label = CLASS_NAMES[int(class_indices[i])]
+            confidence = scores[i]
+
+            # Draw bounding box
+            cv2.rectangle(orig_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+
+            # Put label text
+            label_text = f"{label}: {confidence:.2f}"
+            cv2.putText(orig_image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+    return orig_image
+
+# Process all images in the input folder
+for image_name in os.listdir(input_folder):
+    image_path = os.path.join(input_folder, image_name)
+    output_path = os.path.join(output_folder, image_name)
+
+    # Preprocess
+    image_tensor, orig_image = preprocess_image(image_path)
+
+    # Run inference
+    inputs = {session.get_inputs()[0].name: image_tensor}
+    outputs = session.run(None, inputs)
+
+    # Post-process and save image
+    processed_image = postprocess_output(outputs, orig_image)
+    cv2.imwrite(output_path, processed_image)
+
+    print(f"Processed: {image_name}")
+
+print("Processing complete! Results saved in", output_folder)