Update convert_tflite_2_onnx.py

convert_tflite_2_onnx.py

@@ -40,3 +40,57 @@ output = session.run(None, {input_name: image_data})
 print(output)
 
 
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+
+# Load ONNX model
+onnx_model_path = 'model.onnx'
+session = ort.InferenceSession(onnx_model_path)
+
+# Function to preprocess a single image (resize and normalize)
+def preprocess_image(image_path, input_size=(320, 320)):
+    image = Image.open(image_path).resize(input_size)  # Resize to match model input size
+    image_data = np.array(image).astype('float32')  # Convert to float32
+    image_data = np.expand_dims(image_data, axis=0)  # Add batch dimension (1, height, width, channels)
+    return image_data
+
+# Prepare a batch of images
+image_paths = ['image1.jpg', 'image2.jpg', 'image3.jpg']  # List of image file paths
+batch_size = len(image_paths)
+
+# Preprocess each image and stack them into a batch
+batch_images = np.vstack([preprocess_image(image_path) for image_path in image_paths])
+
+# Check input name from the ONNX model
+input_name = session.get_inputs()[0].name
+
+# Run batch inference
+outputs = session.run(None, {input_name: batch_images})
+
+# Postprocessing: Extract scores, bounding boxes, and labels for each image in the batch
+scores_batch, bboxes_batch, labels_batch = outputs[0], outputs[1], outputs[2]
+
+# Iterate over the batch of results and filter based on score threshold
+score_threshold = 0.5
+
+for i in range(batch_size):
+    scores = scores_batch[i]  # Scores for i-th image
+    bboxes = bboxes_batch[i]  # Bounding boxes for i-th image
+    labels = labels_batch[i]  # Labels for i-th image
+
+    # Filter indices where scores are greater than the threshold
+    valid_indices = np.where(scores > score_threshold)
+
+    # Filter the outputs based on valid indices
+    filtered_scores = scores[valid_indices]
+    filtered_bboxes = bboxes[valid_indices]
+    filtered_labels = labels[valid_indices]
+
+    print(f"Image {i+1}:")
+    print("Filtered Scores:", filtered_scores)
+    print("Filtered Bounding Boxes:", filtered_bboxes)
+    print("Filtered Labels:", filtered_labels)
+    print('---')
+
+