mirror of
https://github.com/StefBuwalda/ProjectIOT.git
synced 2025-10-29 18:59:57 +00:00
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
from ultralytics import YOLO
|
|
from PIL import Image
|
|
import numpy as np
|
|
import easyocr
|
|
|
|
car_model = YOLO("yolov8n.pt")
|
|
plate_model = YOLO("license_plate_detector.pt")
|
|
ocr_reader = easyocr.Reader(["nl"])
|
|
|
|
img = Image.open("test.jpg")
|
|
|
|
results = car_model.predict(source=img)
|
|
|
|
cars: list[tuple[int, tuple[int, int, int, int]]] = []
|
|
|
|
# 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 = ocr_reader.readtext(image=lp_np)
|
|
print(result)
|
|
print(result[0][1]) # type: ignore
|