57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
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
|
|
|
|
|