Add utility functions to Human class for perception

tf_pose/estimator.py

@@ -1,4 +1,5 @@
 import logging
+import math
 
 import slidingwindow as sw
 
@@ -26,6 +27,17 @@
 logger.addHandler(ch)
 
 
+def _round(v):
+    return int(round(v))
+
+
+def _include_part(part_list, part_idx):
+    for part in part_list:
+        if part_idx == part.part_idx:
+            return True, part
+    return False, None
+
+
 class Human:
     """
     body_parts: list of BodyPart
@@ -68,6 +80,158 @@ def part_count(self):
     def get_max_score(self):
         return max([x.score for _, x in self.body_parts.items()])
 
+    def get_face_box(self, img_w, img_h, mode=0):
+        """
+        Get Face box compared to img size (w, h)
+        :param img_w:
+        :param img_h:
+        :param mode:
+        :return:
+        """
+        # SEE : https://github.com/ildoonet/tf-pose-estimation/blob/master/tf_pose/common.py#L13
+        _NOSE = CocoPart.Nose
+        _NECK = CocoPart.Neck
+        _REye = CocoPart.REye
+        _LEye = CocoPart.LEye
+        _REar = CocoPart.REar
+        _LEar = CocoPart.LEar
+
+        _THRESHOLD_PART_CONFIDENCE = 0.2
+        parts = [part for idx, part in self.body_parts.items() if part.score > _THRESHOLD_PART_CONFIDENCE]
+
+        is_nose, part_nose = _include_part(parts, _NOSE)
+        if not is_nose:
+            return None
+
+        size = 0
+        is_neck, part_neck = _include_part(parts, _NECK)
+        if is_neck:
+            size = max(size, img_h * (part_neck.y - part_nose.y) * 0.8)
+
+        is_reye, part_reye = _include_part(parts, _REye)
+        is_leye, part_leye = _include_part(parts, _LEye)
+        if is_reye and is_leye:
+            size = max(size, img_w * (part_reye.x - part_leye.x) * 2.0)
+            size = max(size,
+                       img_w * math.sqrt((part_reye.x - part_leye.x) ** 2 + (part_reye.y - part_leye.y) ** 2) * 2.0)
+
+        if mode == 1:
+            if not is_reye and not is_leye:
+                return None
+
+        is_rear, part_rear = _include_part(parts, _REar)
+        is_lear, part_lear = _include_part(parts, _LEar)
+        if is_rear and is_lear:
+            size = max(size, img_w * (part_rear.x - part_lear.x) * 1.6)
+
+        if size <= 0:
+            return None
+
+        if not is_reye and is_leye:
+            x = part_nose.x * img_w - (size // 3 * 2)
+        elif is_reye and not is_leye:
+            x = part_nose.x * img_w - (size // 3)
+        else:  # is_reye and is_leye:
+            x = part_nose.x * img_w - size // 2
+
+        x2 = x + size
+        if mode == 0:
+            y = part_nose.y * img_h - size // 3
+        else:
+            y = part_nose.y * img_h - _round(size / 2 * 1.2)
+        y2 = y + size
+
+        # fit into the image frame
+        x = max(0, x)
+        y = max(0, y)
+        x2 = min(img_w - x, x2 - x) + x
+        y2 = min(img_h - y, y2 - y) + y
+
+        if _round(x2 - x) == 0.0 or _round(y2 - y) == 0.0:
+            return None
+        if mode == 0:
+            return {"x": _round((x + x2) / 2),
+                    "y": _round((y + y2) / 2),
+                    "w": _round(x2 - x),
+                    "h": _round(y2 - y)}
+        else:
+            return {"x": _round(x),
+                    "y": _round(y),
+                    "w": _round(x2 - x),
+                    "h": _round(y2 - y)}
+
+    def get_upper_body_box(self, img_w, img_h):
+        """
+        Get Upper body box compared to img size (w, h)
+        :param img_w:
+        :param img_h:
+        :return:
+        """
+
+        if not (img_w > 0 and img_h > 0):
+            raise Exception("img size should be positive")
+
+        _NOSE = CocoPart.Nose
+        _NECK = CocoPart.Neck
+        _RSHOULDER = CocoPart.RShoulder
+        _LSHOULDER = CocoPart.LShoulder
+        _THRESHOLD_PART_CONFIDENCE = 0.3
+        parts = [part for idx, part in self.body_parts.items() if part.score > _THRESHOLD_PART_CONFIDENCE]
+        part_coords = [(img_w * part.x, img_h * part.y) for part in parts if
+                       part.part_idx in [0, 1, 2, 5, 8, 11, 14, 15, 16, 17]]
+
+        if len(part_coords) < 5:
+            return None
+
+        # Initial Bounding Box
+        x = min([part[0] for part in part_coords])
+        y = min([part[1] for part in part_coords])
+        x2 = max([part[0] for part in part_coords])
+        y2 = max([part[1] for part in part_coords])
+
+        # # ------ Adjust heuristically +
+        # if face points are detcted, adjust y value
+
+        is_nose, part_nose = _include_part(parts, _NOSE)
+        is_neck, part_neck = _include_part(parts, _NECK)
+        torso_height = 0
+        if is_nose and is_neck:
+            y -= (part_neck.y * img_h - y) * 0.8
+            torso_height = max(0, (part_neck.y - part_nose.y) * img_h * 2.5)
+        #
+        # # by using shoulder position, adjust width
+        is_rshoulder, part_rshoulder = _include_part(parts, _RSHOULDER)
+        is_lshoulder, part_lshoulder = _include_part(parts, _LSHOULDER)
+        if is_rshoulder and is_lshoulder:
+            half_w = x2 - x
+            dx = half_w * 0.15
+            x -= dx
+            x2 += dx
+        elif is_neck:
+            if is_lshoulder and not is_rshoulder:
+                half_w = abs(part_lshoulder.x - part_neck.x) * img_w * 1.15
+                x = min(part_neck.x * img_w - half_w, x)
+                x2 = max(part_neck.x * img_w + half_w, x2)
+            elif not is_lshoulder and is_rshoulder:
+                half_w = abs(part_rshoulder.x - part_neck.x) * img_w * 1.15
+                x = min(part_neck.x * img_w - half_w, x)
+                x2 = max(part_neck.x * img_w + half_w, x2)
+
+        # ------ Adjust heuristically -
+
+        # fit into the image frame
+        x = max(0, x)
+        y = max(0, y)
+        x2 = min(img_w - x, x2 - x) + x
+        y2 = min(img_h - y, y2 - y) + y
+
+        if _round(x2 - x) == 0.0 or _round(y2 - y) == 0.0:
+            return None
+        return {"x": _round((x + x2) / 2),
+                "y": _round((y + y2) / 2),
+                "w": _round(x2 - x),
+                "h": _round(y2 - y)}
+
     def __str__(self):
         return ' '.join([str(x) for x in self.body_parts.values()])
 
@@ -367,7 +531,7 @@ def inference(self, npimg, resize_to_default=True, upsample_size=1.0):
         self.heatMat = heatMat_up[0]
         self.pafMat = pafMat_up[0]
         logger.debug('inference- heatMat=%dx%d pafMat=%dx%d' % (
-        self.heatMat.shape[1], self.heatMat.shape[0], self.pafMat.shape[1], self.pafMat.shape[0]))
+            self.heatMat.shape[1], self.heatMat.shape[0], self.pafMat.shape[1], self.pafMat.shape[0]))
 
         t = time.time()
         humans = PoseEstimator.estimate_paf(peaks, self.heatMat, self.pafMat)