diff --git a/ANPR.py b/ANPR.py index 0b0f470..3c201eb 100644 --- a/ANPR.py +++ b/ANPR.py @@ -1,12 +1,11 @@ from ultralytics import YOLO -from ultralytics.engine.results import Boxes from PIL import Image import numpy as np import easyocr car_model = YOLO("yolov8n.pt") plate_model = YOLO("license_plate_detector.pt") -reader = easyocr.Reader(["nl"]) +ocr_reader = easyocr.Reader(["nl"]) img = Image.open("test.jpg") @@ -42,5 +41,6 @@ for r in results: 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) + result = ocr_reader.readtext(image=lp_np) print(result) + print(result[0][1]) # type: ignore diff --git a/application/__init__.py b/application/__init__.py index 40fec29..2711775 100644 --- a/application/__init__.py +++ b/application/__init__.py @@ -3,6 +3,13 @@ from flask_migrate import Migrate from flask_sqlalchemy import SQLAlchemy from flask_login import LoginManager from io import BytesIO +from ultralytics import YOLO +import easyocr + +# ANPR shiz +car_model = YOLO("yolov8n.pt") +plate_model = YOLO("license_plate_detector.pt") +ocr_reader = easyocr.Reader(["nl"]) # from authlib.integrations.flask_client import OAuth diff --git a/application/api/image_processing.py b/application/api/image_processing.py index c5e3dce..a8556a3 100644 --- a/application/api/image_processing.py +++ b/application/api/image_processing.py @@ -1,19 +1,48 @@ import io +from application import car_model, plate_model, ocr_reader +from PIL import Image +import numpy as np async def process_image(image: bytes) -> str: print("Saving file to memory") image_file = io.BytesIO(image) - print("Processing image") - # anpr_info = await anpr.detect(image_file) # type: ignore + img = Image.open(image_file) - # if anpr_info is None: - # print("Something went wrong with ANPR") - # return "" + results = car_model.predict(source=img) - # if not anpr_info["is_plate"]: - # return "" + cars: list[tuple[int, tuple[int, int, int, int]]] = [] - # print(anpr_info["plate_number"]) - return "" # str(anpr_info["plate_number"]) + # 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) + + return str(result[0][1]) # type: ignore