mirror of
https://github.com/StefBuwalda/ProjectIOT.git
synced 2025-11-01 04:10:00 +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 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
|
||||||
if cls_name == "License_Plate":
|
cars.append((size, (x1, y1, x2, y2)))
|
||||||
conf = float(box2.conf[0])
|
|
||||||
x1, y1, x2, y2 = map(int, box2.xyxy[0])
|
# Get the biggest car box
|
||||||
print(
|
size, corners = max(cars, key=lambda x: x[0])
|
||||||
f"{cls_name} ({conf:.2f}) at [{x1},{y1},{x2},{y2}]"
|
|
||||||
)
|
# Crop biggest car
|
||||||
lp_img = cropped_img.crop((x1, y1, x2, y2))
|
cropped_img = img.crop(corners)
|
||||||
lp_img.save("license_plate.jpg")
|
cropped_img.save("car_crop_pillow.jpg")
|
||||||
lp_np = np.array(lp_img)
|
|
||||||
result3 = reader.readtext(lp_np)
|
# Search for license plates in car box and OCR all
|
||||||
print(result3)
|
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