diff --git a/.gitmodules b/.gitmodules index e5ab9394..e3e0ce23 100644 --- a/.gitmodules +++ b/.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 diff --git a/README.md b/README.md index 2c93a3c4..517bb2c6 100644 --- a/README.md +++ b/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). diff --git a/image_recognition_pytorch/CMakeLists.txt b/image_recognition_pytorch/CMakeLists.txt new file mode 100644 index 00000000..a10619ad --- /dev/null +++ b/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() diff --git a/image_recognition_pytorch/README.md b/image_recognition_pytorch/README.md new file mode 100644 index 00000000..ed285e19 --- /dev/null +++ b/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))] diff --git a/image_recognition_pytorch/doc/face.png b/image_recognition_pytorch/doc/face.png new file mode 100644 index 00000000..d9ae94a3 Binary files /dev/null and b/image_recognition_pytorch/doc/face.png differ diff --git a/image_recognition_pytorch/docs b/image_recognition_pytorch/docs new file mode 160000 index 00000000..6a785e90 --- /dev/null +++ b/image_recognition_pytorch/docs @@ -0,0 +1 @@ +Subproject commit 6a785e90b0039a84f683121fa4742b0c1196acd6 diff --git a/image_recognition_pytorch/package.xml b/image_recognition_pytorch/package.xml new file mode 100644 index 00000000..5b1566c1 --- /dev/null +++ b/image_recognition_pytorch/package.xml @@ -0,0 +1,39 @@ + + + + image_recognition_pytorch + 0.0.1 + The image_recognition_pytorch package + + Loy van Beek + + MIT + + catkin + + python3-setuptools + + diagnostic_updater + image_recognition_msgs + image_recognition_util + python3-numpy + python3-opencv + python3-onnxruntime-pip + python3-pytorch-pip + rospy + + python3-catkin-lint + python3-future + python3-rospkg + + python3-sphinx + python-sphinx-autoapi-pip + python-sphinx-rtd-theme-pip + python3-yaml + + + + + diff --git a/image_recognition_pytorch/rosdoc.yaml b/image_recognition_pytorch/rosdoc.yaml new file mode 100644 index 00000000..d6daeddf --- /dev/null +++ b/image_recognition_pytorch/rosdoc.yaml @@ -0,0 +1,3 @@ +- builder: sphinx + sphinx_root_dir: docs + name: Python API diff --git a/image_recognition_pytorch/scripts/download_model b/image_recognition_pytorch/scripts/download_model new file mode 100755 index 00000000..329484f7 --- /dev/null +++ b/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)) diff --git a/image_recognition_pytorch/scripts/face_properties_node b/image_recognition_pytorch/scripts/face_properties_node new file mode 100755 index 00000000..b28c3310 --- /dev/null +++ b/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) diff --git a/image_recognition_pytorch/scripts/get_face_properties b/image_recognition_pytorch/scripts/get_face_properties new file mode 100755 index 00000000..0b9cee82 --- /dev/null +++ b/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])) diff --git a/image_recognition_pytorch/setup.py b/image_recognition_pytorch/setup.py new file mode 100644 index 00000000..924427d7 --- /dev/null +++ b/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) diff --git a/image_recognition_pytorch/src/image_recognition_pytorch/__init__.py b/image_recognition_pytorch/src/image_recognition_pytorch/__init__.py new file mode 100644 index 00000000..89c696a9 --- /dev/null +++ b/image_recognition_pytorch/src/image_recognition_pytorch/__init__.py @@ -0,0 +1 @@ +from . import age_gender_estimator diff --git a/image_recognition_pytorch/src/image_recognition_pytorch/age_gender_estimator.py b/image_recognition_pytorch/src/image_recognition_pytorch/age_gender_estimator.py new file mode 100644 index 00000000..62cc9d22 --- /dev/null +++ b/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 diff --git a/image_recognition_pytorch/test/assets/age_28_gender_female.jpg b/image_recognition_pytorch/test/assets/age_28_gender_female.jpg new file mode 100644 index 00000000..96feaf3e Binary files /dev/null and b/image_recognition_pytorch/test/assets/age_28_gender_female.jpg differ diff --git a/image_recognition_pytorch/test/assets/age_29_gender_male.jpg b/image_recognition_pytorch/test/assets/age_29_gender_male.jpg new file mode 100644 index 00000000..dab7b7fc Binary files /dev/null and b/image_recognition_pytorch/test/assets/age_29_gender_male.jpg differ diff --git a/image_recognition_pytorch/test/assets/age_33_gender_male.jpg b/image_recognition_pytorch/test/assets/age_33_gender_male.jpg new file mode 100644 index 00000000..946eb978 Binary files /dev/null and b/image_recognition_pytorch/test/assets/age_33_gender_male.jpg differ diff --git a/image_recognition_pytorch/test/run_tests.bash b/image_recognition_pytorch/test/run_tests.bash new file mode 100755 index 00000000..83ee8fff --- /dev/null +++ b/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 )" diff --git a/image_recognition_pytorch/test/test_face_properties.py b/image_recognition_pytorch/test/test_face_properties.py new file mode 100644 index 00000000..5a81e423 --- /dev/null +++ b/image_recognition_pytorch/test/test_face_properties.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python +from __future__ import print_function + +import os +import re +from future.moves.urllib.request import urlretrieve + +import cv2 +import rospkg +from image_recognition_keras.age_gender_estimator import AgeGenderEstimator + + +def test_face_properties(): + local_path = "/tmp/age_gender_weights.hdf5" + + if not os.path.exists(local_path): + http_path = "https://github.com/tue-robotics/image_recognition/releases/download/" \ + "image_recognition_keras_face_properties_weights.28-3.73/" \ + "image_recognition_keras_face_properties_weights.28-3.73.hdf5" + urlretrieve(http_path, local_path) + print("Downloaded weights to {}".format(local_path)) + + def age_is_female_from_asset_name(asset_name): + age_str, gender_str = re.search("age_(\d+)_gender_(\w+)", asset_name).groups() + return int(age_str), gender_str == "female" + + assets_path = os.path.join(rospkg.RosPack().get_path("image_recognition_keras"), 'test/assets') + images_gt = [(cv2.imread(os.path.join(assets_path, asset)), age_is_female_from_asset_name(asset)) + for asset in os.listdir(assets_path)] + + estimations = AgeGenderEstimator(local_path, 64, 16, 8).estimate([image for image, _ in images_gt]) + for (_, (age_gt, is_female_gt)), (age, gender) in zip(images_gt, estimations): + age = int(age) + is_female = gender[0] > 0.5 + assert abs(age - age_gt) < 5 + assert is_female == is_female_gt + + +if __name__ == "__main__": + test_face_properties() diff --git a/image_recognition_rqt/src/image_recognition_rqt/annotation.py b/image_recognition_rqt/src/image_recognition_rqt/annotation.py index 5eeedba8..45ae5a81 100644 --- a/image_recognition_rqt/src/image_recognition_rqt/annotation.py +++ b/image_recognition_rqt/src/image_recognition_rqt/annotation.py @@ -13,8 +13,8 @@ import re import rosservice -from image_widget import ImageWidget -from dialogs import option_dialog, warning_dialog +from .image_widget import ImageWidget +from .dialogs import option_dialog, warning_dialog from image_recognition_msgs.msg import Annotation from image_recognition_util import image_writer diff --git a/image_recognition_rqt/src/image_recognition_rqt/manual.py b/image_recognition_rqt/src/image_recognition_rqt/manual.py index a8f63deb..59f7197b 100644 --- a/image_recognition_rqt/src/image_recognition_rqt/manual.py +++ b/image_recognition_rqt/src/image_recognition_rqt/manual.py @@ -12,8 +12,8 @@ from cv_bridge import CvBridge, CvBridgeError from image_recognition_msgs.msg import CategoryProbability, Recognition -from image_widget import ImageWidget -from dialogs import option_dialog, warning_dialog +from .image_widget import ImageWidget +from .dialogs import option_dialog, warning_dialog from image_recognition_msgs.srv import Recognize, RecognizeResponse import re diff --git a/image_recognition_rqt/src/image_recognition_rqt/test.py b/image_recognition_rqt/src/image_recognition_rqt/test.py index 8e9efa1e..72f70687 100644 --- a/image_recognition_rqt/src/image_recognition_rqt/test.py +++ b/image_recognition_rqt/src/image_recognition_rqt/test.py @@ -2,10 +2,10 @@ import rosservice import rostopic from cv_bridge import CvBridge, CvBridgeError -from dialogs import option_dialog, warning_dialog, info_dialog +from .dialogs import option_dialog, warning_dialog, info_dialog from image_recognition_msgs.msg import CategoryProbability, FaceProperties from image_recognition_msgs.srv import GetFaceProperties, Recognize -from image_widget import ImageWidget +from .image_widget import ImageWidget from python_qt_binding.QtCore import * from python_qt_binding.QtGui import * from python_qt_binding.QtWidgets import *