1import onnxruntime as ort
2import numpy as np
3import cv2
4import easyocr
5
6def preprocess_image(image_path, img_size=(640, 640)):
7 img = cv2.imread(image_path)
8 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
9 img_resized = cv2.resize(img, img_size)
10 img_normalized = img_resized.astype(np.float32) / 255.0
11 img_transposed = np.transpose(img_normalized, (2, 0, 1)) # HWC to CHW
12 img_batch = np.expand_dims(img_transposed, axis=0) # Add batch dimension
13 return img, img_batch
14
15def postprocess_output(output, original_img_shape, img_size=(640, 640), conf_threshold=0.25, iou_threshold=0.7):
16 # output format is (batch_size, 5 + num_classes, num_boxes)
17 predictions = np.squeeze(output).T # Transpose to (num_boxes, 5 + num_classes)
18
19 # Filter out low confidence detections
20 scores = np.max(predictions[:, 4:], axis=1)
21 predictions = predictions[scores > conf_threshold]
22 scores = scores[scores > conf_threshold]
23
24 if len(predictions) == 0:
25 return []
26
27 # Get class IDs and bounding boxes
28 class_ids = np.argmax(predictions[:, 4:], axis=1)
29 boxes = predictions[:, :4]
30
31 # Convert boxes from normalized YOLO format to pixel coordinates
32 # x_center, y_center, width, height
33 x_center, y_center, width, height = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
34 x1 = (x_center - width / 2) * img_size[0]
35 y1 = (y_center - height / 2) * img_size[1]
36 x2 = (x_center + width / 2) * img_size[0]
37 y2 = (y_center + height / 2) * img_size[1]
38
39 # Scale bounding boxes to original image size
40 ratio_w = original_img_shape[1] / img_size[0]
41 ratio_h = original_img_shape[0] / img_size[1]
42
43 x1 *= ratio_w
44 y1 *= ratio_h
45 x2 *= ratio_w
46 y2 *= ratio_h
47
48 bboxes = np.column_stack([x1, y1, x2, y2]).astype(int)
49
50 # Apply Non-Maximum Suppression
51 indices = cv2.dnn.NMSBoxes(bboxes.tolist(), scores.tolist(), conf_threshold, iou_threshold)
52 if len(indices) > 0:
53 indices = indices.flatten()
54 return [{'box': bboxes[i], 'score': scores[i], 'class_id': class_ids[i]} for i in indices]
55 return []
56
57# Load ONNX model
58session = ort.InferenceSession("best.onnx")
59input_name = session.get_inputs()[0].name
60output_name = session.get_outputs()[0].name
61
62# Initialize EasyOCR reader
63reader = easyocr.Reader(['ar', 'en'], gpu=True) # Set gpu=False if no GPU available
64
65# Example usage
66image_path = "path/to/your/image.jpg"
67original_img, preprocessed_img = preprocess_image(image_path)
68
69# Run inference
70outputs = session.run([output_name], {input_name: preprocessed_img})
71
72# Post-process and get detections
73detector_results = postprocess_output(outputs[0], original_img.shape)
74
75if detector_results:
76 for res in detector_results:
77 x1, y1, x2, y2 = res['box']
78 plate_crop = original_img[y1:y2, x1:x2]
79
80 # OCR on the cropped plate
81 ocr_result = reader.readtext(plate_crop, detail=0)
82 plate_text = " ".join(ocr_result)
83 print(f"Detected plate: {plate_text} (Confidence: {res['score']:.2f})")
84
85 # Visualize (optional)
86 img_with_box = original_img.copy()
87 cv2.rectangle(img_with_box, (x1, y1), (x2, y2), (0, 255, 0), 2)
88 cv2.putText(img_with_box, plate_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
89 # plt.imshow(img_with_box)
90 # plt.title(plate_text)
91 # plt.axis('off')
92 # plt.show()
93else:
94 print("No license plate detected.")
95
