File size: 3,692 Bytes
d5c46bf cd66def d5c46bf cd66def d5c46bf cd66def d5c46bf cd66def d5c46bf cd66def d5c46bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
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
# Define unique colors for each class (BGR format for OpenCV)
CLASS_COLORS = {
"person": (255, 0, 0), # Blue
"car": (0, 255, 0), # Green
"truck": (0, 0, 255), # Red
"bicycle": (255, 255, 0), # Cyan
"dog": (255, 0, 255), # Magenta
"cat": (0, 255, 255), # Yellow
}
DEFAULT_COLOR = (200, 200, 200) # Gray for unknown classes
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
# Get class label
class_id = int(class_indices[i])
label = CLASS_NAMES[class_id] if class_id < len(CLASS_NAMES) else "Unknown"
confidence = scores[i]
# Get class-specific color
color = CLASS_COLORS.get(label, DEFAULT_COLOR)
# Draw bounding box
cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
# Put label text
label_text = f"{label}: {confidence:.2f}"
text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
text_x, text_y = x1, y1 - 10
# Draw text background
cv2.rectangle(orig_image, (text_x, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5), color, -1)
# Put label text
cv2.putText(orig_image, label_text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 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)
|