chore: Some minor updates

uniface/detection/retinaface.py

@@ -100,7 +100,8 @@ class RetinaFace:
                 model_path,
                 providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
             )
-            self.input_name = self.session.get_inputs()[0].name
+            self.input_names = self.session.get_inputs()[0].name
+            self.output_names = [x.name for x in self.session.get_outputs()]
             Logger.info(f"Successfully initialized the model from {model_path}")
         except Exception as e:
             Logger.error(f"Failed to load model from '{model_path}': {e}", exc_info=True)
@@ -129,13 +130,13 @@ class RetinaFace:
         Returns:
             Tuple[np.ndarray, np.ndarray]: Raw model outputs.
         """
-        return self.session.run(None, {self.input_name: input_tensor})
+        return self.session.run(self.output_names, {self.input_names: input_tensor})
 
     def detect(
         self,
         image: np.ndarray,
         max_num: Optional[int] = 0,
-        metric: Literal["default", "max"] = "default",
+        metric: Literal["default", "max"] = "max",
         center_weight: Optional[float] = 2.0
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
@@ -159,6 +160,8 @@ class RetinaFace:
                 Shape: (num_detections, 5, 2), where each row contains 5 landmark points (x, y).
         """
 
+        original_height, original_width = image.shape[:2]
+        
         if self.dynamic_size:
             height, width, _ = image.shape
             self._priors = generate_anchors(image_size=(height, width))  # generate anchors for each input image
@@ -180,7 +183,7 @@ class RetinaFace:
             areas = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
 
             # Calculate offsets from image center
-            center = (height // 2, width // 2)
+            center = (original_height // 2, original_width // 2)
             offsets = np.vstack([
                 (detections[:, 0] + detections[:, 2]) / 2 - center[1],
                 (detections[:, 1] + detections[:, 3]) / 2 - center[0]
@@ -241,7 +244,7 @@ class RetinaFace:
 
         # Apply NMS
         detections = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
-        keep = non_max_supression(detections, self.nms_thresh)
+        keep = nms(detections, self.nms_thresh)
         detections, landmarks = detections[keep], landmarks[keep]
 
         # Keep top-k detections
@@ -259,4 +262,4 @@ class RetinaFace:
         landmark_scale = np.array([shape[0], shape[1]] * 5)
         landmarks = landmarks * landmark_scale / resize_factor
 
-        return boxes, landmarks
+        return boxes, landmarks

uniface/detection/scrfd.py

@@ -3,7 +3,7 @@ import cv2
 import numpy as np
 import onnxruntime
 
-from typing import Tuple, Optional
+from typing import Tuple, Optional, List, Literal
 
 # from uniface.logger import Logger
 from .utils import non_max_supression, distance2bbox, distance2kps, resize_image
@@ -42,10 +42,6 @@ class SCRFD:
         self.fmc = 3
         self._feat_stride_fpn = [8, 16, 32]
         self._num_anchors = 2
-        self.use_kps = True
-
-        self.mean = 127.5
-        self.std = 128.0
 
         self.center_cache = {}
         # ---------------------------------
@@ -64,92 +60,124 @@ class SCRFD:
                 providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
             )
             # Get model info
+            self.input_names = self.session.get_inputs()[0].name
             self.output_names = [x.name for x in self.session.get_outputs()]
-            self.input_names = [x.name for x in self.session.get_inputs()]
         except Exception as e:
             print(f"Failed to load the model: {e}")
             raise
 
-    def forward(self, image, threshold):
+    def preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int]]:
+        """Preprocess image for inference.
+
+        Args:
+            image (np.ndarray): Input image
+
+        Returns:
+            Tuple[np.ndarray, Tuple[int, int]]: Preprocessed blob and input size
+        """
+        input_size = tuple(image.shape[0:2][::-1])
+
+        image = image.astype(np.float32)
+        image = (image - 127.5) / 127.5
+        image = image.transpose(2, 0, 1)  # HWC to CHW
+        image = np.expand_dims(image, axis=0)
+
+        return image, input_size
+
+    def inference(self, input_tensor: np.ndarray) -> List[np.ndarray]:
+        """Perform model inference on the preprocessed image tensor.
+
+        Args:
+            input_tensor (np.ndarray): Preprocessed input tensor.
+
+        Returns:
+            Tuple[np.ndarray, np.ndarray]: Raw model outputs.
+        """
+        return self.session.run(self.output_names, {self.input_names: input_tensor})
+
+    def postprocess(self, outputs: List[np.ndarray], image_dim: Tuple[int, int]):
         scores_list = []
         bboxes_list = []
         kpss_list = []
-        input_size = tuple(image.shape[0:2][::-1])
 
-        blob = cv2.dnn.blobFromImage(
-            image,
-            1.0 / self.std,
-            input_size,
-            (self.mean, self.mean, self.mean),
-            swapRB=True
-        )
-        outputs = self.session.run(self.output_names, {self.input_names[0]: blob})
-
-        input_height = blob.shape[2]
-        input_width = blob.shape[3]
+        input_height, input_width = image_dim
 
         fmc = self.fmc
         for idx, stride in enumerate(self._feat_stride_fpn):
             scores = outputs[idx]
-            bbox_preds = outputs[idx + fmc]
-            bbox_preds = bbox_preds * stride
-            if self.use_kps:
-                kps_preds = outputs[idx + fmc * 2] * stride
+            bbox_preds = outputs[fmc + idx] * stride
+            kps_preds = outputs[2*fmc + idx] * stride
 
-            height = input_height // stride
-            width = input_width // stride
-            key = (height, width, stride)
-            if key in self.center_cache:
-                anchor_centers = self.center_cache[key]
+            # Generate anchors
+            fm_height = input_height // stride
+            fm_width = input_width // stride
+            cache_key = (fm_height, fm_width, stride)
+
+            if cache_key in self.center_cache:
+                anchor_centers = self.center_cache[cache_key]
             else:
-                anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
-                anchor_centers = (anchor_centers * stride).reshape((-1, 2))
-                if self._num_anchors > 1:
-                    anchor_centers = np.stack([anchor_centers] * self._num_anchors, axis=1).reshape((-1, 2))
-                if len(self.center_cache) < 100:
-                    self.center_cache[key] = anchor_centers
+                y, x = np.mgrid[:fm_height, :fm_width]
+                anchor_centers = np.stack((x, y), axis=-1).astype(np.float32)
+                anchor_centers = (anchor_centers * stride).reshape(-1, 2)
+
+                if self._num_anchors > 1:
+                    anchor_centers = np.tile(anchor_centers[:, None, :], (1, self._num_anchors, 1)).reshape(-1, 2)
+
+                if len(self.center_cache) < 100:
+                    self.center_cache[cache_key] = anchor_centers
+
+            pos_indices = np.where(scores >= self.conf_thres)[0]
+            if len(pos_indices) == 0:
+                continue
+
+            bboxes = distance2bbox(anchor_centers, bbox_preds)[pos_indices]
+            scores_selected = scores[pos_indices]
+            scores_list.append(scores_selected)
+            bboxes_list.append(bboxes)
+
+            kpss = distance2kps(anchor_centers, kps_preds)
+            kpss = kpss.reshape((kpss.shape[0], -1, 2))
+            kpss_list.append(kpss[pos_indices])
 
-            pos_inds = np.where(scores >= threshold)[0]
-            bboxes = distance2bbox(anchor_centers, bbox_preds)
-            pos_scores = scores[pos_inds]
-            pos_bboxes = bboxes[pos_inds]
-            scores_list.append(pos_scores)
-            bboxes_list.append(pos_bboxes)
-            if self.use_kps:
-                kpss = distance2kps(anchor_centers, kps_preds)
-                kpss = kpss.reshape((kpss.shape[0], -1, 2))
-                pos_kpss = kpss[pos_inds]
-                kpss_list.append(pos_kpss)
         return scores_list, bboxes_list, kpss_list
 
-    def detect(self, image, max_num=1, metric="max", center_weight: Optional[float] = 2) -> Tuple[np.ndarray, np.ndarray]:
-        original_height, original_width = image.shape[:2]
-        image, resize_factor = resize_image(image, target_shape=self.input_size)
-        
+    def detect(
+        self,
+        image: np.ndarray,
+        max_num: Optional[int] = 0,
+        metric: Literal["default", "max"] = "max",
+        center_weight: Optional[float] = 2
+    ) -> Tuple[np.ndarray, np.ndarray]:
 
-        scores_list, bboxes_list, kpss_list = self.forward(image, self.conf_thres)
+        original_height, original_width = image.shape[:2]
+
+        image, resize_factor = resize_image(image, target_shape=self.input_size)
+
+        image_tensor, _ = self.preprocess(image)
+
+        outputs = self.inference(image_tensor)
+
+        scores_list, bboxes_list, kpss_list = self.postprocess(outputs, image.shape[:2])
 
         scores = np.vstack(scores_list)
         scores_ravel = scores.ravel()
         order = scores_ravel.argsort()[::-1]
         bboxes = np.vstack(bboxes_list) / resize_factor
 
-        if self.use_kps:
-            kpss = np.vstack(kpss_list) / resize_factor
+        kpss = np.vstack(kpss_list) / resize_factor
 
         pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False)
         pre_det = pre_det[order, :]
         keep = non_max_supression(pre_det, threshold=self.iou_thres)
         det = pre_det[keep, :]
-        if self.use_kps:
-            kpss = kpss[order, :, :]
-            kpss = kpss[keep, :, :]
-        else:
-            kpss = None
+
+        kpss = kpss[order, :, :]
+        kpss = kpss[keep, :, :]
+
         if 0 < max_num < det.shape[0]:
             area = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1])
             center = (original_height // 2, original_width // 2)
-            
+
             offsets = np.vstack(
                 [
                     (det[:, 0] + det[:, 2]) / 2 - center[1],
@@ -161,24 +189,26 @@ class SCRFD:
                 values = area
             else:
                 values = area - offset_dist_squared * center_weight  # some extra weight on the centering
-                
+
             sorted_indices = np.argsort(values)[::-1][:max_num]
             det = det[sorted_indices]
-            if kpss is not None:
-                kpss = kpss[sorted_indices]
-            
+            kpss = kpss[sorted_indices]
+
         return det, kpss
 
+
 def draw_bbox(frame, bbox, color=(0, 255, 0), thickness=2):
     x1, y1, x2, y2 = bbox[:4].astype(np.int32)
     cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness)
     score = bbox[4]
     cv2.putText(frame, f"{score:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
 
+
 def draw_keypoints(frame, points, color=(0, 0, 255), radius=2):
     for (x, y) in points.astype(np.int32):
         cv2.circle(frame, (x, y), radius, color, -1)
 
+
 if __name__ == "__main__":
     detector = SCRFD(model_path="det_10g.onnx")
     cap = cv2.VideoCapture(0)
@@ -208,4 +238,4 @@ if __name__ == "__main__":
             break
 
     cap.release()
-    cv2.destroyAllWindows()
+    cv2.destroyAllWindows()

uniface/retinaface.py

@@ -100,7 +100,8 @@ class RetinaFace:
                 model_path,
                 providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
             )
-            self.input_name = self.session.get_inputs()[0].name
+            self.input_names = self.session.get_inputs()[0].name
+            self.output_names = [x.name for x in self.session.get_outputs()]
             Logger.info(f"Successfully initialized the model from {model_path}")
         except Exception as e:
             Logger.error(f"Failed to load model from '{model_path}': {e}", exc_info=True)
@@ -129,13 +130,13 @@ class RetinaFace:
         Returns:
             Tuple[np.ndarray, np.ndarray]: Raw model outputs.
         """
-        return self.session.run(None, {self.input_name: input_tensor})
+        return self.session.run(self.output_names, {self.input_names: input_tensor})
 
     def detect(
         self,
         image: np.ndarray,
         max_num: Optional[int] = 0,
-        metric: Literal["default", "max"] = "default",
+        metric: Literal["default", "max"] = "max",
         center_weight: Optional[float] = 2.0
     ) -> Tuple[np.ndarray, np.ndarray]:
         """