Merge pull request #113 from StefBuwalda/ANPR

Nieuw ANPR systeem

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

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)

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"

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

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