Standardize on PyTorch for machine learning models

.gitmodules

@@ -28,3 +28,6 @@
 [submodule "image_recognition_util/docs"]
 	path = image_recognition_util/docs
 	url = https://github.com/tue-robotics/tue_documentation_python.git
+[submodule "image_recognition_pytorch/docs"]
+	path = image_recognition_pytorch/docs
+	url = https://github.com/tue-robotics/tue_documentation_python.git

README.md

@@ -17,6 +17,7 @@ Package | Build status Xenial Kinetic x64 | Description
 [image_recognition_msgs](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_msgs) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_msgs__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_msgs__ubuntu_xenial__source/1/) | Interface definition for image recognition
 [image_recognition_openface](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_openface) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_openface__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_openface__ubuntu_xenial__source/1/) | ROS wrapper for Openface (https://github.com/cmusatyalab/openface) to detect and recognize faces in images.
 [image_recognition_openpose](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_openpose) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_openpose__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_openpose_ubuntu_xenial__source/1/) | ROS wrapper for Openpose (https://github.com/CMU-Perceptual-Computing-Lab/) for getting poses of 2D images.
+[image_recognition_pytorch](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_pytorch) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_pytorch__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_pytorch_ubuntu_xenial__source/1/) | ROS wrapper around a PyTorch model for (https://github.com/Nebula4869/PyTorch-gender-age-estimation) for getting age & gender estimations on faces
 [image_recognition_rqt](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_rqt) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_rqt__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_rqt__ubuntu_xenial__source/1/) | RQT tools with helpers testing this interface and training/labeling data.
 [image_recognition_skybiometry](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_skybiometry) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_skybiometry__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_skybiometry_ubuntu_xenial__source/1/) | ROS wrapper for Skybiometry (https://skybiometry.com/) for getting face properties of a detected face, e.g. age estimation, gender estimation etc.
 [image_recognition_tensorflow](https://github.com/tue-robotics/image_recognition/tree/master/image_recognition_tensorflow) | [![Build Status](http://build.ros.org/job/Ksrc_uX__image_recognition_tensorflow__ubuntu_xenial__source/1//badge/icon)](http://build.ros.org/job/Ksrc_uX__image_recognition_tensorflow__ubuntu_xenial__source/1/) | Object recognition with use of Tensorflow. The user can retrain the top layers of a neural network to perform classification with its own dataset as described in [this tutorial](https://www.tensorflow.org/versions/r0.11/how_tos/image_retraining/index.html).

image_recognition_pytorch/CMakeLists.txt

@@ -0,0 +1,25 @@
+cmake_minimum_required(VERSION 3.0.2)
+project(image_recognition_pytorch)
+
+find_package(catkin REQUIRED)
+
+catkin_python_setup()
+
+catkin_package()
+
+install(PROGRAMS
+  scripts/face_properties_node
+  scripts/get_face_properties
+  DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
+)
+
+if (CATKIN_ENABLE_TESTING)
+    # Test catkin lint
+    find_program(CATKIN_LINT catkin_lint REQUIRED)
+    execute_process(COMMAND "${CATKIN_LINT}" "-q" "-W2" "${CMAKE_SOURCE_DIR}" RESULT_VARIABLE lint_result)
+    if(NOT ${lint_result} EQUAL 0)
+      message(FATAL_ERROR "catkin_lint failed")
+    endif()
+
+    catkin_add_nosetests(test)
+endif()

image_recognition_pytorch/README.md

@@ -0,0 +1,50 @@
+# Image recognition pytorch
+
+Image recognition (age and gender estimation of a face) with use of PyTorch.
+
+## Installation
+
+See https://github.com/tue-robotics/image_recognition for installation instructions.
+
+## ROS Node (face_properties_node)
+
+Age and gender estimation
+```
+rosrun image_recognition_pytorch face_properties_node _weights_file_path:=[path_to_model]
+```
+
+Run the image_recognition_rqt test gui (https://github.com/tue-robotics/image_recognition_rqt)
+
+    rosrun image_recognition_rqt test_gui
+
+Configure the service you want to call with the gear-wheel in the top-right corner of the screen. If everything is set-up, draw a rectangle in the image around a face:
+
+![Wide ResNet](doc/wide_resnet_test.png)
+
+## Scripts
+
+### Download model
+
+Download weights from github.
+
+```
+usage: download_model [-h] [--model_path MODEL_PATH]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --model_path MODEL_PATH
+```
+
+### Get face properties (get_face_properties)
+
+Get the classification result of an input image:
+
+```
+rosrun image_recognition_pytorch get_face_properties `rospack find image_recognition_pytorch`/doc/face.png
+```
+
+![Example](doc/face.png)
+
+Output:
+
+    [(50.5418073660112, array([0.5845756 , 0.41542447], dtype=float32))]

image_recognition_pytorch/package.xml

@@ -0,0 +1,39 @@
+<?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_pytorch</name>
+  <version>0.0.1</version>
+  <description>The image_recognition_pytorch package</description>
+
+  <maintainer email="[email protected]">Loy van Beek</maintainer>
+
+  <license>MIT</license>
+
+  <buildtool_depend>catkin</buildtool_depend>
+
+  <buildtool_depend>python3-setuptools</buildtool_depend>
+
+  <exec_depend>diagnostic_updater</exec_depend>
+  <exec_depend>image_recognition_msgs</exec_depend>
+  <exec_depend>image_recognition_util</exec_depend>
+  <exec_depend>python3-numpy</exec_depend>
+  <exec_depend>python3-opencv</exec_depend>
+  <exec_depend>python3-onnxruntime-pip</exec_depend>
+  <exec_depend>python3-pytorch-pip</exec_depend>
+  <exec_depend>rospy</exec_depend>
+
+  <test_depend>python3-catkin-lint</test_depend>
+  <test_depend>python3-future</test_depend>
+  <test_depend>python3-rospkg</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_pytorch/rosdoc.yaml

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

image_recognition_pytorch/scripts/download_model

@@ -0,0 +1,23 @@
+#!/usr/bin/env python
+from __future__ import print_function
+import os
+import urllib.request
+
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model_path', default=os.path.expanduser('~/data/pytorch_models'))
+args = parser.parse_args()
+
+os.system('mkdir -p {}'.format(args.model_path))
+local_path = os.path.join(args.model_path, 'best-epoch47-0.9314.onnx')
+
+if not os.path.exists(local_path):
+    # TODO: Clone this for us
+    http_path = "https://github.com/Nebula4869/PyTorch-gender-age-estimation/raw/" \
+                "038331d26fc1fbf24d00365d0eb9d0e5e828dda6/models-2020-11-20-14-37/best-epoch47-0.9314.onnx"
+    print("Downloading model to {} ...".format(local_path))
+    urllib.request.urlretrieve(http_path, local_path)
+    print("Model downloaded: {}".format(local_path))
+else:
+    print("Model already downloaded: {}".format(local_path))

image_recognition_pytorch/scripts/face_properties_node

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+import os
+import sys
+
+import diagnostic_updater
+import rospy
+from cv_bridge import CvBridge, CvBridgeError
+from image_recognition_pytorch.age_gender_estimator import AgeGenderEstimator
+from image_recognition_msgs.msg import FaceProperties
+from image_recognition_msgs.srv import GetFaceProperties
+from image_recognition_util import image_writer
+
+
+class PytorchFaceProperties:
+    def __init__(self, weights_file_path, img_size, depth, width, save_images_folder, use_gpu):
+        """
+        ROS node that wraps the PyTorch age gender estimator
+        """
+        self._bridge = CvBridge()
+        self._properties_srv = rospy.Service('get_face_properties', GetFaceProperties, self._get_face_properties_srv)
+        self._estimator = AgeGenderEstimator(weights_file_path, img_size, depth, width, use_gpu)
+
+        if save_images_folder:
+            self._save_images_folder = os.path.expanduser(save_images_folder)
+            if not os.path.exists(self._save_images_folder):
+                os.makedirs(self._save_images_folder)
+        else:
+            self._save_images_folder = None
+
+        rospy.loginfo("PytorchFaceProperties node initialized:")
+        rospy.loginfo(" - weights_file_path=%s", weights_file_path)
+        rospy.loginfo(" - img_size=%s", img_size)
+        rospy.loginfo(" - depth=%s", depth)
+        rospy.loginfo(" - width=%s", width)
+        rospy.loginfo(" - save_images_folder=%s", save_images_folder)
+        rospy.loginfo(" - use_gpu=%s", use_gpu)
+
+    def _get_face_properties_srv(self, req):
+        """
+        Callback when the GetFaceProperties service is called
+
+        :param req: Input images
+        :return: properties
+        """
+        # Convert to opencv images
+        try:
+            bgr_images = [self._bridge.imgmsg_to_cv2(image, "bgr8") for image in req.face_image_array]
+        except CvBridgeError as e:
+            raise Exception("Could not convert image to opencv image: %s" % str(e))
+
+        rospy.loginfo("Estimating the age and gender of %d incoming images ...", len(bgr_images))
+        estimations = self._estimator.estimate(bgr_images)
+        rospy.loginfo("Done")
+
+        face_properties_array = []
+        for (age, gender_prob) in estimations:
+            gender, gender_confidence = (FaceProperties.FEMALE, gender_prob[0]) if gender_prob[0] > 0.5 else (FaceProperties.MALE, gender_prob[1])
+
+            face_properties_array.append(FaceProperties(
+                age=int(age),
+                gender=gender,
+                gender_confidence=gender_confidence
+            ))
+
+        # Store images if specified
+        if self._save_images_folder:
+            def _get_label(p):
+                return "age_%d_gender_%s" % (p.age, "male" if p.gender == FaceProperties.MALE else "female")
+
+            image_writer.write_estimations(self._save_images_folder, bgr_images,
+                                           [_get_label(p) for p in face_properties_array],
+                                           suffix="_face_properties")
+
+        # Service response
+        return {"properties_array": face_properties_array}
+
+
+if __name__ == '__main__':
+    rospy.init_node("face_properties")
+
+    try:
+        default_weights_path = os.path.expanduser('~/data/pytorch_models/best-epoch47-0.9314.onnx')
+        weights_file_path = rospy.get_param("~weights_file_path", default_weights_path)
+        img_size = rospy.get_param("~image_size", 64)
+        depth = rospy.get_param("~depth", 16)
+        width = rospy.get_param("~width", 8)
+        save_images = rospy.get_param("~save_images", True)
+        use_gpu = rospy.get_param("~use_gpu", False)
+
+        save_images_folder = None
+        if save_images:
+            save_images_folder = rospy.get_param("~save_images_folder", "/tmp/image_recognition_pytorch")
+    except KeyError as e:
+        rospy.logerr("Parameter %s not found" % e)
+        sys.exit(1)
+
+    try:
+        PytorchFaceProperties(weights_file_path, img_size, depth, width, save_images_folder, use_gpu)
+        updater = diagnostic_updater.Updater()
+        updater.setHardwareID("none")
+        updater.add(diagnostic_updater.Heartbeat())
+        rospy.Timer(rospy.Duration(1), lambda event: updater.force_update())
+        rospy.spin()
+    except Exception as e:
+        rospy.logfatal(e)

image_recognition_pytorch/scripts/get_face_properties

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+from __future__ import print_function
+import argparse
+from image_recognition_pytorch.age_gender_estimator import AgeGenderEstimator
+import cv2
+import os
+
+# Assign description to the help doc
+parser = argparse.ArgumentParser(description='Get face properties using PyTorch')
+
+# Add arguments
+parser.add_argument('image', type=str, help='Image')
+parser.add_argument('--weights-path', type=str, help='Path to the weights of the WideResnet model',
+                    default=os.path.expanduser('~/data/pytorch_models/best-epoch47-0.9314.onnx'))
+parser.add_argument('--image-size', type=int, help='Size of the input image', default=64)
+parser.add_argument('--depth', type=int, help='Depth of the network', default=16)
+parser.add_argument('--width', type=int, help='Width of the network', default=8)
+
+args = parser.parse_args()
+
+# Read the image
+img = cv2.imread(args.image)
+
+estimator = AgeGenderEstimator(args.weights_path, args.image_size, args.depth, args.width)
+
+print(estimator.estimate([img]))

image_recognition_pytorch/setup.py

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

image_recognition_pytorch/src/image_recognition_pytorch/__init__.py

@@ -0,0 +1 @@
+from . import age_gender_estimator

image_recognition_pytorch/src/image_recognition_pytorch/age_gender_estimator.py

@@ -0,0 +1,61 @@
+import cv2
+import numpy as np
+import os.path
+
+import onnxruntime
+
+GENDER_DICT = {0: 'male', 1: 'female'}
+
+
+class AgeGenderEstimator(object):
+    def __init__(self, weights_file_path, img_size=64, depth=16, width=8, use_gpu=False):
+        """
+        Estimate the age and gender of the incoming image
+
+        :param weights_file_path: path to a pre-trained network in onnx format
+        """
+        weights_file_path = os.path.expanduser(weights_file_path)
+
+        if not os.path.isfile(weights_file_path):
+            raise IOError("Weights file {}, no such file ..".format(weights_file_path))
+
+        self._model = None
+        self._weights_file_path = weights_file_path
+        self._img_size = img_size
+        self._depth = depth
+        self._width = width
+        self._use_gpu = use_gpu
+
+    def estimate(self, np_images):
+        """
+        Estimate the age and gender of the face on the image
+
+        :param np_images a numpy array of BGR images of faces of which the gender and the age has to be estimated
+            This is assumed to be segmented/cropped already!
+        :returns List of estimated age and gender score ([female, male]) tuples
+        """
+
+        # Model should be constructed in same thread as the inference
+        if self._model is None:
+            providers = ['CPUExecutionProvider']
+            if self._use_gpu:
+                providers.append(
+                    ('CUDAExecutionProvider', {
+                        'device_id': 0,
+                        'arena_extend_strategy': 'kNextPowerOfTwo',
+                        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
+                        'cudnn_conv_algo_search': 'EXHAUSTIVE',
+                        'do_copy_in_default_stream': True,
+                    })),
+
+            self._model = onnxruntime.InferenceSession(self._weights_file_path, providers=providers)
+
+        results = []
+        for np_image in np_images:
+            inputs = np.transpose(cv2.resize(np_image, (64, 64)), (2, 0, 1))
+            inputs = np.expand_dims(inputs, 0).astype(np.float32) / 255.
+            predictions = self._model.run(['output'], input_feed={'input': inputs})[0][0]
+            #           age              p(male)          p(female)
+            results += [(predictions[2], (predictions[0], predictions[1]))]
+
+        return results

image_recognition_pytorch/test/run_tests.bash

@@ -0,0 +1,2 @@
+#!/bin/bash
+nosetests -vv "$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"