Only OCR the biggest bounding box

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

53
ANPR.py
View File

@@ -1,4 +1,5 @@
from ultralytics import YOLO from ultralytics import YOLO
from ultralytics.engine.results import Boxes
from PIL import Image from PIL import Image
import numpy as np import numpy as np
import easyocr import easyocr
@@ -11,31 +12,35 @@ img = Image.open("test.jpg")
results = car_model.predict(source=img) results = car_model.predict(source=img)
for r in results: cars: list[tuple[int, tuple[int, int, int, int]]] = []
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}]")
# Crop Image # Filter out the cars and calculate box size
cropped_img = img.crop((x1, y1, x2, y2)) for r in results:
cropped_img.save("car_crop_pillow.jpg") if r.boxes:
results2 = plate_model.predict(source=cropped_img) for box in r.boxes:
for x in results2: cls_name = r.names[int(box.cls[0])]
for box2 in x.boxes: if cls_name == "car":
cls_id = int(box2.cls[0]) x1, y1, x2, y2 = map(int, box.xyxy[0])
cls_name = x.names[cls_id] 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": if cls_name == "License_Plate":
conf = float(box2.conf[0]) x1, y1, x2, y2 = map(int, box.xyxy[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 = cropped_img.crop((x1, y1, x2, y2))
lp_img.save("license_plate.jpg") lp_img.save("license_plate.jpg")
lp_np = np.array(lp_img) lp_np = np.array(object=lp_img)
result3 = reader.readtext(lp_np) result = reader.readtext(image=lp_np)
print(result3) print(result)