diff --git a/ANPR.py b/ANPR.py new file mode 100644 index 0000000..3c201eb --- /dev/null +++ b/ANPR.py @@ -0,0 +1,46 @@ +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 diff --git a/app.py b/app.py index 2cd25f2..e3bc026 100644 --- a/app.py +++ b/app.py @@ -16,4 +16,4 @@ def home(): if __name__ == "__main__": - app.run(host="localhost", port=2222, debug=True) + app.run(host="0.0.0.0", port=2222, debug=True) diff --git a/application/__init__.py b/application/__init__.py index 7f6faba..2711775 100644 --- a/application/__init__.py +++ b/application/__init__.py @@ -2,9 +2,21 @@ from flask import Flask 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 +# Memory for last_image +last_image = BytesIO() + # Web Server app = Flask(__name__) app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///data.sqlite" 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 diff --git a/application/api/views.py b/application/api/views.py index 2505323..e1e4088 100644 --- a/application/api/views.py +++ b/application/api/views.py @@ -1,9 +1,11 @@ from flask import Blueprint, request, jsonify from application import db +import application from application.dashboard.models import AllowedPlate, LoggedItem from application.api.image_processing import process_image from datetime import datetime import asyncio +import io api_blueprint = Blueprint("api", __name__, template_folder="templates") @@ -12,6 +14,8 @@ api_blueprint = Blueprint("api", __name__, template_folder="templates") @api_blueprint.route("/", methods=["POST"]) def data(): data = request.data + application.last_image.seek(0) + application.last_image = io.BytesIO(request.data) np = asyncio.run(process_image(image=data)) # Check if the found plate is allowed to exit diff --git a/car_crop_pillow.jpg b/car_crop_pillow.jpg new file mode 100644 index 0000000..447682a Binary files /dev/null and b/car_crop_pillow.jpg differ diff --git a/license_plate.jpg b/license_plate.jpg new file mode 100644 index 0000000..c6835a9 Binary files /dev/null and b/license_plate.jpg differ diff --git a/license_plate_detector.pt b/license_plate_detector.pt new file mode 100644 index 0000000..a0ebd6b Binary files /dev/null and b/license_plate_detector.pt differ diff --git a/test.jpg b/test.jpg new file mode 100644 index 0000000..e188371 Binary files /dev/null and b/test.jpg differ diff --git a/yolov8n.pt b/yolov8n.pt new file mode 100644 index 0000000..0db4ca4 Binary files /dev/null and b/yolov8n.pt differ