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)