mirror of
https://github.com/StefBuwalda/ProjectIOT.git
synced 2025-10-29 18:59:57 +00:00
Only OCR the biggest bounding box
This commit is contained in:
59
ANPR.py
59
ANPR.py
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user