BlurMe/ml/element_detection.py

import os
from typing import List, Tuple

import torch

from ultralytics import YOLO

DIR_PATH = os.path.dirname(os.path.realpath(__file__))
FACES_PATH = os.path.join(DIR_PATH, "assets/weights/yolov8n-face.pt")
PLATES_PATH = os.path.join(DIR_PATH, "assets/weights/yolov8-plate.pt")
FACES_MODEL = YOLO(FACES_PATH)
PLATES_MODEL = YOLO(PLATES_PATH)
CONF_THRESH = 0.3
IOU_THRESH = 0.5


class BoundBox:
    def __init__(self, x1, y1, x2, y2, object=None):
        self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
        self.selected = True
        if object not in ["face", "plate"]:
            raise ValueError("object must be either 'face' or 'plate'")
        self.object = object

    def select(self):
        self.selected = True

    def unselect(self):
        self.selected = False

    def get_params(self) -> Tuple[int, int, int, int]:
        return self.x1, self.y1, self.x2, self.y2


def detect(image_path: str) -> List[BoundBox]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    faces = FACES_MODEL.predict(
        source=image_path, conf=CONF_THRESH, iou=IOU_THRESH, device=device
    )
    faces = faces[0].cpu().numpy().boxes
    plates = PLATES_MODEL.predict(
        source=image_path, conf=CONF_THRESH, iou=IOU_THRESH, device=device
    )
    plates = plates[0].cpu().numpy().boxes
    bounding_boxes = []
    for boxes, tag in zip([faces, plates], ["face", "plate"]):
        for box in boxes:
            xyxyn = box.xyxy[0]
            x1 = int(xyxyn[0])
            y1 = int(xyxyn[1])
            x2 = int(xyxyn[2])
            y2 = int(xyxyn[3])
            bounding_boxes.append(BoundBox(x1, y1, x2, y2, tag))
    return bounding_boxes