Merge pull request #18 from Ar-Ray-code/features/yolox_ros_0.3.0

Features/yolox ros 0.3.0

Author: Ar-Ray-code

README.md

@@ -32,7 +32,7 @@ Check [GitHub Wiki](https://github.com/Ar-Ray-code/YOLOX-ROS/wiki/YOLOX-ROS---Ra
 - OpenCV 4
 - Python 3.8 (Ubuntu 20.04 Default)
 - PyTorch >= v1.7
-- [YOLOX v0.2.0](https://github.com/Megvii-BaseDetection/YOLOX)
+- [YOLOX v0.3.0](https://github.com/Megvii-BaseDetection/YOLOX)
 - [bbox_ex_msgs](https://github.com/Ar-Ray-code/bbox_ex_msgs)
 
 ## Installation
@@ -50,15 +50,15 @@ git clone https://github.com/Ar-Ray-code/yolox_ros.git --recursive
 ### STEP 2 : YOLOX Installation (yolox_rps_py)
 
 ```bash
-bash YOLOX-ROS/yolox_ros_py/install_yolox_py.bash
+pip3 install yolox
 ```
 
 ### STEP 3 : Install YOLOX-ROS
 
 ```bash
 source /opt/ros/foxy/setup.bash
 sudo apt install ros-foxy-v4l2-camera
-colcon build --symlink-install # weights files will be installed automatically.
+colcon build --symlink-install # weights (YOLOX-Nano) files will be installed automatically.
 ```
 
 ### (Step 3) Using CUDA
@@ -78,8 +78,13 @@ Connect your web camera.
 ```bash
 source /opt/ros/foxy/setup.bash
 source ~/ros2_ws/install/local_setup.bash
-ros2 launch yolox_ros_py yolox_s_cpu.launch.py
-# ros2 launch yolox_ros_py yolox_s.launch.py # <- GPU
+ros2 launch yolox_ros_py yolox_nano_cpu.launch.py
+# ros2 launch yolox_ros_py yolox_nano.launch.py # <- GPU
+# ros2 launch yolox_ros_py yolox_nano_onnx.launch.py # <- ONNXRuntime
+
+# OpenVINO ---
+# source /opt/intel/openvino_2021/bin/setupvars.sh
+# ros2 launch yolox_ros_py yolox_nano_openvino.launch.py # <- OpenVINO
 ```
 
 </details>

yolox_ros_py/exps/yolox_nano.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.25
+        self.input_size = (416, 416)
+        self.random_size = (10, 20)
+        self.mosaic_scale = (0.5, 1.5)
+        self.test_size = (416, 416)
+        self.mosaic_prob = 0.5
+        self.enable_mixup = False
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self, sublinear=False):
+
+        def init_yolo(M):
+            for m in M.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eps = 1e-3
+                    m.momentum = 0.03
+        if "model" not in self.__dict__:
+            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+            in_channels = [256, 512, 1024]
+            # NANO model use depthwise = True, which is main difference.
+            backbone = YOLOPAFPN(
+                self.depth, self.width, in_channels=in_channels,
+                act=self.act, depthwise=True,
+            )
+            head = YOLOXHead(
+                self.num_classes, self.width, in_channels=in_channels,
+                act=self.act, depthwise=True
+            )
+            self.model = YOLOX(backbone, head)
+
+        self.model.apply(init_yolo)
+        self.model.head.initialize_biases(1e-2)
+        return self.model

yolox_ros_py/install_yolox_py.bash

@@ -19,13 +19,13 @@ def generate_launch_description():
         parameters=[
             {"image_size/width": 640},
             {"image_size/height": 480},
-            {"yolox_exp_py" : yolox_ros_share_dir+'/yolox_s.py'},
+            {"yolox_exp_py" : yolox_ros_share_dir+'/yolox_nano.py'},
             {"device" : 'gpu'},
             {"fp16" : True},
             {"fuse" : False},
             {"legacy" : False},
             {"trt" : False},
-            {"ckpt" : yolox_ros_share_dir+"/yolox_s.pth"},
+            {"ckpt" : yolox_ros_share_dir+"/yolox_nano.pth"},
             {"conf" : 0.3},
             {"threshold" : 0.65},
             {"resize" : 640},

yolox_ros_py/launch/yolox_s.launch.py renamed to yolox_ros_py/launch/yolox_nano.launch.py

@@ -19,13 +19,13 @@ def generate_launch_description():
         parameters=[
             {"image_size/width": 640},
             {"image_size/height": 480},
-            {"yolox_exp_py" : yolox_ros_share_dir+'/yolox_s.py'},
+            {"yolox_exp_py" : yolox_ros_share_dir+'/yolox_nano.py'},
             {"device" : 'cpu'},
             {"fp16" : True},
             {"fuse" : False},
             {"legacy" : False},
             {"trt" : False},
-            {"ckpt" : yolox_ros_share_dir+"/yolox_s.pth"},
+            {"ckpt" : yolox_ros_share_dir+"/yolox_nano.pth"},
             {"conf" : 0.3},
             {"threshold" : 0.65},
             {"resize" : 640},

yolox_ros_py/launch/yolox_s_cpu.launch.py renamed to yolox_ros_py/launch/yolox_nano_cpu.launch.py

@@ -0,0 +1,42 @@
+import launch
+import launch_ros.actions
+from launch.actions import IncludeLaunchDescription
+from ament_index_python.packages import get_package_share_directory
+from launch.launch_description_sources import PythonLaunchDescriptionSource
+
+from urllib.request import urlretrieve
+import os
+
+def generate_launch_description():
+    yolox_ros_share_dir = get_package_share_directory('yolox_ros_py')
+
+    webcam = launch_ros.actions.Node(
+        package="v4l2_camera", executable="v4l2_camera_node",
+        parameters=[
+            {"image_size": [640,480]},
+        ],
+    )
+    yolox_openvino = launch_ros.actions.Node(
+        package="yolox_ros_py", executable="yolox_onnx",output="screen",
+        parameters=[
+            {"image_size/width": 640},
+            {"image_size/height": 480},
+            
+            {"input_shape/width": 416},
+            {"input_shape/height": 416},
+
+            {"with_p6" : False},
+            {"model_path" : yolox_ros_share_dir+"/yolox_nano.onnx"},
+            {"conf" : 0.3},
+        ],
+    )
+
+    rqt_graph = launch_ros.actions.Node(
+        package="rqt_graph", executable="rqt_graph",
+    )
+
+    return launch.LaunchDescription([
+        webcam,
+        yolox_openvino,
+        # rqt_graph
+    ])

yolox_ros_py/launch/yolox_nano_onnx.launch.py

@@ -22,7 +22,7 @@ def generate_launch_description():
             {"image_size/width": 640},
             {"image_size/height": 480},
             {"device" : 'CPU'},
-            {"model_path" : yolox_ros_share_dir+"/yolox_s.xml"},
+            {"model_path" : yolox_ros_share_dir+"/yolox_nano.onnx"},
             {"conf" : 0.3},
         ],
     )

yolox_ros_py/launch/yolox_s_openvino.launch.py renamed to yolox_ros_py/launch/yolox_nano_openvino.launch.py

@@ -2,7 +2,7 @@
 <?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
 <package format="3">
   <name>yolox_ros_py</name>
-  <version>0.2.1</version>
+  <version>0.3.0</version>
   <description>The yolox_ros_py package</description>
   <maintainer email="[email protected]">Ar-Ray-code</maintainer>
   <license>Apache License, Version 2.0</license>

yolox_ros_py/package.xml

@@ -7,47 +7,67 @@
 package_name = 'yolox_ros_py'
 
 
-YOLOX_S_WEIGHTS = 'yolox_s.pth'
-YOLOX_M_WEIGHTS = 'yolox_m.pth'
-YOLOX_L_WEIGHTS = 'yolox_l.pth'
-YOLOX_X_WEIGHTS = 'yolox_x.pth'
-YOLOX_DARKNET53_WEIGHTS = 'yolox_darknet53.pth'
-YOLOX_NANO_WEIGHTS = 'yolox_nano.pth'
-YOLOX_TINY_WEIGHTS = 'yolox_tiny.pth'
+YOLOX_S = 'yolox_s'
+YOLOX_M = 'yolox_m'
+YOLOX_L = 'yolox_l'
+YOLOX_X = 'yolox_x'
+YOLOX_DARKNET53 = 'yolox_darknet53'
+YOLOX_NANO = 'yolox_nano'
+YOLOX_TINY = 'yolox_tiny'
+
+PTH = '.pth'
+ONNX = '.onnx'
 
 YOLOX_VERSION = '0.1.1rc0'
 
 BASE_LINK = 'https://github.com/Megvii-BaseDetection/YOLOX/releases/download/'+YOLOX_VERSION+'/'
-YOLOX_S_WEIGHTS_URL = BASE_LINK + YOLOX_S_WEIGHTS
-YOLOX_M_WEIGHTS_URL = BASE_LINK + YOLOX_M_WEIGHTS
-YOLOX_L_WEIGHTS_URL = BASE_LINK + YOLOX_L_WEIGHTS
-YOLOX_X_WEIGHTS_URL = BASE_LINK + YOLOX_X_WEIGHTS
-YOLOX_DARKNET53_WEIGHTS_URL = BASE_LINK + YOLOX_DARKNET53_WEIGHTS
-YOLOX_NANO_WEIGHTS_URL = BASE_LINK + YOLOX_NANO_WEIGHTS
-YOLOX_TINY_WEIGHTS_URL = BASE_LINK + YOLOX_TINY_WEIGHTS
+# .pth
+YOLOX_S_URL = BASE_LINK + YOLOX_S + PTH
+YOLOX_M_URL = BASE_LINK + YOLOX_M + PTH
+YOLOX_L_URL = BASE_LINK + YOLOX_L + PTH
+YOLOX_X_URL = BASE_LINK + YOLOX_X + PTH
+YOLOX_X_URL = BASE_LINK + YOLOX_DARKNET53 + PTH
+YOLOX_NANO_URL = BASE_LINK + YOLOX_NANO + PTH
+YOLOX_TINY_URL = BASE_LINK + YOLOX_TINY + PTH
+# .onnx
+YOLOX_S_ONNX_URL = BASE_LINK + YOLOX_S + ONNX
+YOLOX_M_ONNX_URL = BASE_LINK + YOLOX_M + ONNX
+YOLOX_L_ONNX_URL = BASE_LINK + YOLOX_L + ONNX
+YOLOX_X_ONNX_URL = BASE_LINK + YOLOX_X + ONNX
+YOLOX_X_ONNX_URL = BASE_LINK + YOLOX_DARKNET53 + ONNX
+YOLOX_NANO_ONNX_URL = BASE_LINK + YOLOX_NANO + ONNX
+YOLOX_TINY_ONNX_URL = BASE_LINK + YOLOX_TINY + ONNX
 
 BASE_PATH = os.getcwd() + '/../weights/'
 os.makedirs(BASE_PATH, exist_ok=True)
-YOLOX_S_WEIGHTS_PATH = BASE_PATH + YOLOX_S_WEIGHTS
-YOLOX_M_WEIGHTS_PATH = BASE_PATH + YOLOX_M_WEIGHTS
-YOLOX_L_WEIGHTS_PATH = BASE_PATH + YOLOX_L_WEIGHTS
-YOLOX_X_WEIGHTS_PATH = BASE_PATH + YOLOX_X_WEIGHTS
-YOLOX_DARKNET53_WEIGHTS_PATH = BASE_PATH + YOLOX_DARKNET53_WEIGHTS
-YOLOX_NANO_WEIGHTS_PATH = BASE_PATH + YOLOX_NANO_WEIGHTS
-YOLOX_TINY_WEIGHTS_PATH = BASE_PATH + YOLOX_TINY_WEIGHTS
-
-# Download YOLOX-S Weights
-if not os.path.exists(YOLOX_S_WEIGHTS_PATH):
-    urlretrieve(YOLOX_S_WEIGHTS_URL, YOLOX_S_WEIGHTS_PATH)
+# .pth
+YOLOX_S_PATH = BASE_PATH + YOLOX_S + PTH
+YOLOX_M_PATH = BASE_PATH + YOLOX_M + PTH
+YOLOX_L_PATH = BASE_PATH + YOLOX_L + PTH
+YOLOX_X_PATH = BASE_PATH + YOLOX_X + PTH
+YOLOX_DARKNET53_PATH = BASE_PATH + YOLOX_DARKNET53 + PTH
+YOLOX_NANO_PATH = BASE_PATH + YOLOX_NANO + PTH
+YOLOX_TINY_PATH = BASE_PATH + YOLOX_TINY + PTH
+# .onnx
+YOLOX_S_ONNX_PATH = BASE_PATH + YOLOX_S + ONNX
+YOLOX_M_ONNX_PATH = BASE_PATH + YOLOX_M + ONNX
+YOLOX_L_ONNX_PATH = BASE_PATH + YOLOX_L + ONNX
+YOLOX_X_ONNX_PATH = BASE_PATH + YOLOX_X + ONNX
+YOLOX_DARKNET53_ONNX_PATH = BASE_PATH + YOLOX_DARKNET53 + ONNX
+YOLOX_NANO_ONNX_PATH = BASE_PATH + YOLOX_NANO + ONNX
+YOLOX_TINY_ONNX_PATH = BASE_PATH + YOLOX_TINY + ONNX
 
-# Download examples
-# if not os.path.exists(YOLOX_TINY_WEIGHTS_PATH):
-#     urlretrieve(YOLOX_TINY_WEIGHTS_URL, YOLOX_TINY_WEIGHTS_PATH)
+# Download YOLOX-NANO Weights
+if not os.path.exists(YOLOX_NANO_PATH):
+    urlretrieve(YOLOX_NANO_URL, YOLOX_NANO_PATH)
+# Download YOLOX-NANO ONNX
+if not os.path.exists(YOLOX_NANO_ONNX_PATH):
+    urlretrieve(YOLOX_NANO_ONNX_URL, YOLOX_NANO_ONNX_PATH)
 
 
 setup(
     name=package_name,
-    version='0.2.0',
+    version='0.3.0',
     packages=[package_name],
     data_files=[
         ('share/ament_index/resource_index/packages',
@@ -56,7 +76,8 @@
         (os.path.join('share', package_name), glob('./launch/*.launch.py')),
         (os.path.join('share', package_name), glob('../weights/*.pth')),
         (os.path.join('share', package_name), glob('../weights/openvino/*')),
-        (os.path.join('share', package_name), glob('./YOLOX/exps/default/*.py')),
+        (os.path.join('share', package_name), glob('../weights/onnx/*')),
+        (os.path.join('share', package_name), glob('./exps/*.py')),
     ],
     install_requires=['setuptools'],
     zip_safe=True,
@@ -71,6 +92,7 @@
         'console_scripts': [
             'yolox_ros = '+package_name+'.yolox_ros:ros_main',
             'yolox_openvino = '+package_name+'.yolox_openvino:ros_main',
+            'yolox_onnx = '+package_name+'.yolox_onnx:ros_main',
         ],
     },
 )