detection dataset: Adjust scale parameter of head and face detections

Author: JeansBoussier

lada/lib/box_utils.py

@@ -1,6 +1,27 @@
-from lada.lib import Box
+import random
+
+from lada.lib import Box, Image
 
 def box_overlap(box1: Box, box2: Box):
     y1min, x1min, y1max, x1max = box1
     y2min, x2min, y2max, x2max = box2
-    return x1min < x2max and x2min < x1max and y1min < y2max and y2min < y1max
+    return x1min < x2max and x2min < x1max and y1min < y2max and y2min < y1max
+
+def scale_box(img: Image, box: Box, mask_scale=1.0) -> Box:
+    img_h, img_w = img.shape[:2]
+    s = mask_scale - 1.0
+    t, l, b, r = box
+    w, h = r - l + 1, b - t + 1
+    t -= h * s
+    b += h * s
+    l -= w * s
+    r += w * s
+    t = max(0, t)
+    b = min(img_h - 1, b)
+    l = max(0, l)
+    r = min(img_w - 1, r)
+    return int(t), int(l), int(b), int(r)
+
+def random_scale_box(img: Image, box: Box, scale_range=(1.0, 1.5)) -> Box:
+    scale = random.uniform(scale_range[0], scale_range[1])
+    return scale_box(img, box, scale)

lada/lib/face_detector.py

@@ -4,19 +4,9 @@
 import cv2
 import numpy as np
 from lada.lib import Mask, Box, Image, Detection, Detections, DETECTION_CLASSES
-from lada.lib import mask_utils
+from lada.lib import box_utils
 from lada.centerface.centerface import CenterFace
 
-def scale_box(box, mask_scale=1.0) -> Box:
-    s = mask_scale - 1.0
-    t, l, b, r = box
-    w, h = r - l + 1, b - t + 1
-    t -= h * s
-    b += h * s
-    l -= w * s
-    r += w * s
-    return int(t), int(l), int(b), int(r)
-
 def convert_to_boxes(dets) -> list[Box]:
     boxes = []
     for i, det in enumerate(dets):
@@ -50,17 +40,15 @@ def create_mask(frame: Image, box: Box) -> Mask:
 
     return mask
 
-def get_nsfw_frame(dets: list[Box], frame, random_extend_masks: bool, mask_scale: float) -> Detections | None:
+def get_nsfw_frame(dets: list[Box], frame: Image, random_extend_masks: bool) -> Detections | None:
     if len(dets) == 0:
         return None
     detections = []
     for box in dets:
-        box = scale_box(box, mask_scale)
-        mask = create_mask(frame, box)
 
         if random_extend_masks:
-            mask = mask_utils.apply_random_mask_extensions(mask)
-            box = mask_utils.get_box(mask)
+            box = box_utils.random_scale_box(frame, box, scale_range=(1.2, 1.5))
+        mask = create_mask(frame, box)
 
         t, l, b, r = box
         width, height = r - l + 1, b - t + 1
@@ -72,14 +60,13 @@ def get_nsfw_frame(dets: list[Box], frame, random_extend_masks: bool, mask_scale
     return Detections(frame, detections)
 
 class FaceDetector:
-    def __init__(self, model: CenterFace, random_extend_masks=False, conf=0.2, mask_scale=1.3):
+    def __init__(self, model: CenterFace, random_extend_masks=False, conf=0.2):
         self.model = model
         self.random_extend_masks = random_extend_masks
         self.conf = conf
-        self.mask_scale = mask_scale
 
     def detect(self, file_path: str) -> Detections | None:
         image = cv2.imread(file_path, cv2.IMREAD_COLOR_RGB)
         dets, _ = self.model(image, threshold=self.conf)
         dets_boxes = convert_to_boxes(dets)
-        return get_nsfw_frame(dets_boxes, cv2.cvtColor(image, cv2.COLOR_RGB2BGR), random_extend_masks=self.random_extend_masks, mask_scale=self.mask_scale)
+        return get_nsfw_frame(dets_boxes, cv2.cvtColor(image, cv2.COLOR_RGB2BGR), random_extend_masks=self.random_extend_masks)

lada/lib/head_detector.py

@@ -4,7 +4,7 @@
 import cv2
 import numpy as np
 from lada.lib import Mask, Box, Image, Detections, Detection, DETECTION_CLASSES
-from lada.lib import mask_utils
+from lada.lib import box_utils
 from lada.bpjdet.inference import inference
 
 def _create_mask(frame: Image, box: Box) -> Mask:
@@ -37,11 +37,9 @@ def _get_detection(dets: list[Box], frame, random_extend_masks: bool) -> Detecti
         return None
     detections = []
     for box in dets:
-        mask = _create_mask(frame, box)
-
         if random_extend_masks:
-            mask = mask_utils.apply_random_mask_extensions(mask)
-            box = mask_utils.get_box(mask)
+            box = box_utils.random_scale_box(frame, box, scale_range=(0.9, 1.2))
+        mask = _create_mask(frame, box)
 
         t, l, b, r = box
         width, height = r - l + 1, b - t + 1
@@ -52,7 +50,7 @@ def _get_detection(dets: list[Box], frame, random_extend_masks: bool) -> Detecti
     return Detections(frame, detections)
 
 class HeadDetector:
-    def __init__(self, model, data, conf_thres,  iou_thres, imgz=1024, random_extend_masks=False):
+    def __init__(self, model, data, conf_thres,  iou_thres, imgz=1536, random_extend_masks=False):
         self.model = model
         self.data = data
         self.random_extend_masks = random_extend_masks