BLUR-15 Zmiana detekcji twarzy na detekcję obiektów. Dodanie klasy dla ramek wokół obiektów. Dodanie wizualizacji.

ml/element_detection.py

@@ -0,0 +1,69 @@
+import os
+from typing import List, Tuple
+
+import torch
+from PIL import Image, ImageDraw, ImageFont
+
+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):
+        self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
+        self.selected = True
+
+    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 in [faces, plates]:
+        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))
+    return bounding_boxes
+
+
+def show_image_with_boxes(
+    in_image_path: str, bounding_boxes: List[BoundBox], out_image_path: str = None
+):
+    img = Image.open(in_image_path)
+    draw = ImageDraw.Draw(img)
+    font_path = DIR_PATH + "/assets/fonts/arial.ttf"
+    font = ImageFont.truetype(font_path, 25)
+    for i, box in enumerate(bounding_boxes):
+        draw.rectangle(box.get_params(), outline="red", width=2, fill=None)
+        draw.text((box.x1 + 5, box.y1 + 5), str(i+1), fill="red", font=font)
+    if not out_image_path:
+        out_image_path = (
+            in_image_path.split(".")[0] + "_out." + in_image_path.split(".")[1]
+        )
+    img.save(out_image_path)

ml/face_detection.py

@@ -1,29 +0,0 @@
-import os
-from typing import List
-
-import torch
-
-from ultralytics import YOLO
-
-WEIGHTS_PATH = "assets/weights/yolov8n-face.pt"
-DIR_PATH = os.path.dirname(os.path.realpath(__file__))
-WEIGHTS_PATH = os.path.join(DIR_PATH, WEIGHTS_PATH)
-MODEL = YOLO(WEIGHTS_PATH)
-CONF_THRESH = 0.01
-IOU_THRESH = 0.5
-
-
-# TODO: currently detect_faces accepts a image path, but it can be changed to accept images in memory
-def detect_faces(image_path: str) -> List[tuple[int, int, int, int]]:
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    results = MODEL.predict(source=image_path, conf=CONF_THRESH, iou=IOU_THRESH, device=device)
-    face_boxes = []
-    result = results[0].cpu().numpy()
-    for box in result.boxes:
-        xyxyn = box.xyxy[0]
-        x1 = int(xyxyn[0])
-        y1 = int(xyxyn[1])
-        x2 = int(xyxyn[2])
-        y2 = int(xyxyn[3])
-        face_boxes.append((x1, y1, x2, y2))
-    return face_boxes