Only OCR the biggest bounding box

This commit is contained in:
2025-06-12 11:34:42 +02:00
parent 15d9d33ba8
commit d422268a03

59
ANPR.py
View File

@@ -1,4 +1,5 @@
from ultralytics import YOLO
from ultralytics.engine.results import Boxes
from PIL import Image
import numpy as np
import easyocr
@@ -11,31 +12,35 @@ img = Image.open("test.jpg")
results = car_model.predict(source=img)
for r in results:
for box in r.boxes:
cls_id = int(box.cls[0])
cls_name = r.names[cls_id]
if cls_name == "car":
conf = float(box.conf[0])
x1, y1, x2, y2 = map(int, box.xyxy[0])
print(f"{cls_name} ({conf:.2f}) at [{x1},{y1},{x2},{y2}]")
cars: list[tuple[int, tuple[int, int, int, int]]] = []
# Crop Image
cropped_img = img.crop((x1, y1, x2, y2))
cropped_img.save("car_crop_pillow.jpg")
results2 = plate_model.predict(source=cropped_img)
for x in results2:
for box2 in x.boxes:
cls_id = int(box2.cls[0])
cls_name = x.names[cls_id]
if cls_name == "License_Plate":
conf = float(box2.conf[0])
x1, y1, x2, y2 = map(int, box2.xyxy[0])
print(
f"{cls_name} ({conf:.2f}) at [{x1},{y1},{x2},{y2}]"
)
lp_img = cropped_img.crop((x1, y1, x2, y2))
lp_img.save("license_plate.jpg")
lp_np = np.array(lp_img)
result3 = reader.readtext(lp_np)
print(result3)
# Filter out the cars and calculate box size
for r in results:
if r.boxes:
for box in r.boxes:
cls_name = r.names[int(box.cls[0])]
if cls_name == "car":
x1, y1, x2, y2 = map(int, box.xyxy[0])
size = (x2 - x1) ** 2 + (y2 - y1) ** 2
cars.append((size, (x1, y1, x2, y2)))
# Get the biggest car box
size, corners = max(cars, key=lambda x: x[0])
# Crop biggest car
cropped_img = img.crop(corners)
cropped_img.save("car_crop_pillow.jpg")
# Search for license plates in car box and OCR all
results = plate_model.predict(source=cropped_img)
for r in results:
if r.boxes:
for box in r.boxes:
cls_name = r.names[int(box.cls[0])]
if cls_name == "License_Plate":
x1, y1, x2, y2 = map(int, box.xyxy[0])
lp_img = cropped_img.crop((x1, y1, x2, y2))
lp_img.save("license_plate.jpg")
lp_np = np.array(object=lp_img)
result = reader.readtext(image=lp_np)
print(result)