(pose_estimation) new pkg (#215)

.gitmodules

@@ -34,3 +34,7 @@
 [submodule "image_recognition_util/docs"]
 	path = image_recognition_util/docs
 	url = https://github.com/tue-robotics/tue_documentation_python.git
+[submodule "image_recognition_pose_estimation/docs"]
+	path = image_recognition_pose_estimation/docs
+	url = https://github.com/tue-robotics/tue_documentation_python.git
+	branch = master

image_recognition/package.xml

@@ -18,7 +18,7 @@
   <exec_depend>image_recognition_footwear</exec_depend>
   <exec_depend>image_recognition_msgs</exec_depend>
   <exec_depend>image_recognition_openface</exec_depend>
-  <exec_depend>image_recognition_openpose</exec_depend>
+  <exec_depend>image_recognition_pose_estimation</exec_depend>
   <exec_depend>image_recognition_rqt</exec_depend>
   <exec_depend>image_recognition_skybiometry</exec_depend>
   <exec_depend>image_recognition_tensorflow</exec_depend>

image_recognition_pose_estimation/CMakeLists.txt

@@ -0,0 +1,18 @@
+cmake_minimum_required(VERSION 3.5)
+project(image_recognition_pose_estimation)
+
+find_package(catkin REQUIRED COMPONENTS)
+
+catkin_python_setup()
+
+catkin_package()
+
+install(PROGRAMS
+  scripts/detect_poses
+  DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
+)
+
+if(CATKIN_ENABLE_TESTING)
+  find_package(catkin_lint_cmake REQUIRED)
+  catkin_add_catkin_lint_test()
+endif()

image_recognition_pose_estimation/README.md

@@ -0,0 +1,3 @@
+# image_recognition_pose_estimation
+
+This package contains the pose estimation library and node.

image_recognition_pose_estimation/package.xml

@@ -0,0 +1,35 @@
+<?xml version="1.0"?>
+<?xml-model
+  href="http://download.ros.org/schema/package_format3.xsd"
+  schematypens="http://www.w3.org/2001/XMLSchema"?>
+<package format="3">
+  <name>image_recognition_pose_estimation</name>
+  <version>0.0.1</version>
+  <description>The image_recognition_pose_estimation package</description>
+
+  <author email="[email protected]">Matthijs van der Burgh</author>
+
+  <maintainer email="[email protected]">Matthijs van der Burgh</maintainer>
+
+  <license>MIT</license>
+
+  <buildtool_depend>catkin</buildtool_depend>
+
+  <buildtool_depend>python3-setuptools</buildtool_depend>
+
+  <exec_depend>python3-numpy</exec_depend>
+  <exec_depend>python3-opencv</exec_depend>
+  <exec_depend>python3-ultralytics-pip</exec_depend>
+  <exec_depend>rospy</exec_depend>
+
+  <test_depend>catkin_lint_cmake</test_depend>
+
+  <doc_depend>python3-sphinx</doc_depend>
+  <doc_depend>python-sphinx-autoapi-pip</doc_depend>
+  <doc_depend>python-sphinx-rtd-theme-pip</doc_depend>
+  <doc_depend>python3-yaml</doc_depend>
+
+  <export>
+    <rosdoc config="rosdoc.yaml" />
+  </export>
+</package>

image_recognition_pose_estimation/rosdoc.yaml

@@ -0,0 +1,3 @@
+- builder: sphinx
+  sphinx_root_dir: docs
+  name: Python API

image_recognition_pose_estimation/scripts/detect_poses

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+
+import argparse
+import logging
+import os
+import sys
+
+import cv2
+
+from image_recognition_pose_estimation.yolo_pose_wrapper import YoloPoseWrapper
+
+parser = argparse.ArgumentParser(
+    description="Detect poses in an image", formatter_class=argparse.ArgumentDefaultsHelpFormatter
+)
+parser.add_argument("--pose_model", help="What pose model to use", default="yolov8n-pose.pt")
+parser.add_argument(
+    "--verbose",
+    action="store_true",
+    help="If True, enables verbose output during the model's operations. Defaults to False.",
+)
+
+mode_parser = parser.add_subparsers(help="Mode")
+image_parser = mode_parser.add_parser("image", help="Use image mode")
+image_parser.set_defaults(mode="image")
+cam_parser = mode_parser.add_parser("cam", help="Use cam mode")
+cam_parser.set_defaults(mode="cam")
+
+# Image specific arguments
+image_parser.add_argument("image", help="Input image")
+
+args = parser.parse_args()
+
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+
+wrapper = YoloPoseWrapper(args.pose_model, args.verbose)
+
+if args.mode == "image":
+    # Read the image
+    image = cv2.imread(args.image)
+    recognitions, overlayed_image = wrapper.detect_poses(image)
+
+    logging.info(recognitions)
+    cv2.imshow("overlayed_image", overlayed_image)
+
+    cv2.waitKey()
+
+elif args.mode == "cam":
+    cap = cv2.VideoCapture(0)
+    while True:
+        ret, img = cap.read()
+        recognitions, overlayed_image = wrapper.detect_poses(img)
+        cv2.imshow("overlayed_image", overlayed_image)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            break

image_recognition_pose_estimation/setup.py

@@ -0,0 +1,6 @@
+from catkin_pkg.python_setup import generate_distutils_setup
+from setuptools import setup
+
+d = generate_distutils_setup(packages=["image_recognition_pose_estimation"], package_dir={"": "src"})
+
+setup(**d)

image_recognition_pose_estimation/src/image_recognition_pose_estimation/yolo_pose_wrapper.py

@@ -0,0 +1,45 @@
+import logging
+import os
+import re
+import sys
+from typing import List
+
+import numpy as np
+from image_recognition_msgs.msg import CategoricalDistribution, CategoryProbability, Recognition
+from sensor_msgs.msg import RegionOfInterest
+from ultralytics import YOLO
+from ultralytics.engine.results import Results
+
+YOLO_POSE_PATTERN = re.compile(r"^yolov8[nsmlx]-pose(-p6)?.pt$")
+
+
+class YoloPoseWrapper:
+    def __init__(self, model_name: str = "yolov8n-pose.pt", verbose: bool = False):
+        if not YOLO_POSE_PATTERN.match(model_name):
+            raise ValueError(f"Model name '{model_name}' does not match pattern '{YOLO_POSE_PATTERN.pattern}'")
+
+        self._model = YOLO(model=model_name, task="pose", verbose=verbose)
+
+    def detect_poses(self, image: np.ndarray):
+        # Detect poses
+        results: List[Results] = self._model.__call__(image)
+        recognitions = []
+
+        if keypoints is not None and len(keypoints.shape) == 3:  # If no detections, keypoints will be None
+            num_persons, num_bodyparts, _ = keypoints.shape
+            for person_id in range(0, num_persons):
+                for body_part_id in range(0, num_bodyparts):
+                    body_part = self._model["body_parts"][body_part_id]
+                    x, y, probability = keypoints[person_id][body_part_id]
+                    if probability > 0:
+                        recognitions.append(
+                            Recognition(
+                                group_id=person_id,
+                                roi=RegionOfInterest(width=1, height=1, x_offset=int(x), y_offset=int(y)),
+                                categorical_distribution=CategoricalDistribution(
+                                    probabilities=[CategoryProbability(label=body_part, probability=float(probability))]
+                                ),
+                            )
+                        )
+
+        return recognitions, overlayed_image

image_recognition_pose_estimation/test/test_import.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+
+import unittest
+
+
+class TestImport(unittest.TestCase):
+    def test_import(self):
+        """
+        If no exception is raised, this test will succeed
+        """
+        import image_recognition_pose_estimation
+
+
+if __name__ == "__main__":
+    unittest.main()