Moved ANPR system to image processing

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

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

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