Skip to content

Commit

Permalink
OAK-1-MAX non stereo / non spatial YOLO example node (#262)
Browse files Browse the repository at this point in the history 
* Add support for OAK-1-MAX, OAK-1-Lite, and OAK-1 non-spatial cameras

* Add OAK-1-Lite to description

* "True" != "true"

---------

Co-authored-by: MAVProxyUser <[email protected]>
  • Loading branch information
@MAVProxyUser
MAVProxyUser and MAVProxyUser authored May 4, 2023
1 parent d4666b9 commit 29c2906
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 69 deletions.
17 changes: 10 additions & 7 deletions depthai_examples/launch/yolov4_publisher.launch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def generate_launch_description():
tf_prefix = LaunchConfiguration('tf_prefix', default = 'oak')
base_frame = LaunchConfiguration('base_frame', default = 'oak-d_frame')
parent_frame = LaunchConfiguration('parent_frame', default = 'oak-d-base-frame')
spatial_camera = LaunchConfiguration('spatial_camera', default = True)

cam_pos_x = LaunchConfiguration('cam_pos_x', default = '0.0')
cam_pos_y = LaunchConfiguration('cam_pos_y', default = '0.0')
Expand All @@ -34,6 +35,7 @@ def generate_launch_description():
camera_param_uri = LaunchConfiguration('camera_param_uri', default = 'package://depthai_examples/params/camera')
sync_nn = LaunchConfiguration('sync_nn', default = True)
subpixel = LaunchConfiguration('subpixel', default = True)

nnName = LaunchConfiguration('nnName', default = "x")
resourceBaseFolder = LaunchConfiguration('resourceBaseFolder', default = default_resources_path)
confidence = LaunchConfiguration('confidence', default = 200)
Expand All @@ -43,7 +45,7 @@ def generate_launch_description():
declare_camera_model_cmd = DeclareLaunchArgument(
'camera_model',
default_value=camera_model,
description='The model of the camera. Using a wrong camera model can disable camera features. Valid models: `OAK-D, OAK-D-LITE`.')
description='The model of the camera. Using a wrong camera model can disable camera features. Set spatial_camera:=False for `OAK-1, OAK-1-LITE, OAK-1-MAX`.')

declare_tf_prefix_cmd = DeclareLaunchArgument(
'tf_prefix',
Expand Down Expand Up @@ -143,7 +145,6 @@ def generate_launch_description():
'cam_roll' : cam_roll,
'cam_pitch' : cam_pitch,
'cam_yaw' : cam_yaw}.items())

yolov4_spatial_node = launch_ros.actions.Node(
package='depthai_examples', executable='yolov4_spatial_node',
output='screen',
Expand All @@ -153,6 +154,7 @@ def generate_launch_description():
{'nnName': nnName},
{'resourceBaseFolder': resourceBaseFolder},
{'monoResolution': monoResolution}])
{'spatial_camera': spatial_camera}])

rviz_node = launch_ros.actions.Node(
package='rviz2', executable='rviz2', output='screen',
Expand All @@ -173,16 +175,17 @@ def generate_launch_description():
ld.add_action(declare_yaw_cmd)

ld.add_action(declare_camera_param_uri_cmd)
ld.add_action(declare_sync_nn_cmd)
ld.add_action(declare_subpixel_cmd)
ld.add_action(declare_nnName_cmd)
ld.add_action(declare_resourceBaseFolder_cmd)
ld.add_action(declare_confidence_cmd)
ld.add_action(declare_lrCheckTresh_cmd)
ld.add_action(declare_monoResolution_cmd)
ld.add_action(declare_sync_nn_cmd)
ld.add_action(urdf_launch)

if spatial_camera == True:
ld.add_action(declare_subpixel_cmd)
ld.add_action(declare_lrCheckTresh_cmd)
ld.add_action(declare_monoResolution_cmd)
ld.add_action(yolov4_spatial_node)
ld.add_action(urdf_launch)

return ld

165 changes: 103 additions & 62 deletions depthai_examples/src/yolov4_spatial_publisher.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include "camera_info_manager/camera_info_manager.hpp"
#include "depthai_bridge/BridgePublisher.hpp"
#include "depthai_bridge/ImageConverter.hpp"
#include "depthai_bridge/ImgDetectionConverter.hpp"
#include "depthai_bridge/SpatialDetectionConverter.hpp"
#include "depthai_ros_msgs/msg/spatial_detection_array.hpp"
#include "rclcpp/executors.hpp"
Expand All @@ -18,6 +19,7 @@
#include "depthai/pipeline/node/ColorCamera.hpp"
#include "depthai/pipeline/node/MonoCamera.hpp"
#include "depthai/pipeline/node/SpatialDetectionNetwork.hpp"
#include "depthai/pipeline/node/DetectionNetwork.hpp"
#include "depthai/pipeline/node/StereoDepth.hpp"
#include "depthai/pipeline/node/XLinkOut.hpp"

Expand All @@ -32,83 +34,121 @@ const std::vector<std::string> label_map = {
"laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
"refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"};

dai::Pipeline createPipeline(bool syncNN, bool subpixel, std::string nnPath, int confidence, int LRchecktresh, std::string resolution) {
dai::Pipeline createPipeline(bool spatial_camera, bool syncNN, bool subpixel, std::string nnPath, int confidence, int LRchecktresh, std::string resolution) {
dai::Pipeline pipeline;
dai::node::MonoCamera::Properties::SensorResolution monoResolution;
auto colorCam = pipeline.create<dai::node::ColorCamera>();
auto camRgb = pipeline.create<dai::node::ColorCamera>(); // non spatial add in
auto spatialDetectionNetwork = pipeline.create<dai::node::YoloSpatialDetectionNetwork>();
auto detectionNetwork = pipeline.create<dai::node::YoloDetectionNetwork>();
auto monoLeft = pipeline.create<dai::node::MonoCamera>();
auto monoRight = pipeline.create<dai::node::MonoCamera>();
auto stereo = pipeline.create<dai::node::StereoDepth>();


// create xlink connections
auto xoutRgb = pipeline.create<dai::node::XLinkOut>();
auto xoutDepth = pipeline.create<dai::node::XLinkOut>();
auto xoutNN = pipeline.create<dai::node::XLinkOut>();
auto nnOut = pipeline.create<dai::node::XLinkOut>(); // non spatial add in

xoutRgb->setStreamName("preview");
xoutNN->setStreamName("detections");
xoutDepth->setStreamName("depth");

colorCam->setPreviewSize(416, 416);
colorCam->setResolution(dai::ColorCameraProperties::SensorResolution::THE_1080_P);
colorCam->setInterleaved(false);
colorCam->setColorOrder(dai::ColorCameraProperties::ColorOrder::BGR);

if(resolution == "720p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_720_P;
} else if(resolution == "400p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_400_P;
} else if(resolution == "800p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_800_P;
} else if(resolution == "480p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_480_P;
} else {
RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "Invalid parameter. -> monoResolution: %s", resolution.c_str());
throw std::runtime_error("Invalid mono camera resolution.");
if(spatial_camera == true) {
xoutDepth->setStreamName("depth");

colorCam->setPreviewSize(416, 416);
colorCam->setResolution(dai::ColorCameraProperties::SensorResolution::THE_1080_P);
colorCam->setInterleaved(false);
colorCam->setColorOrder(dai::ColorCameraProperties::ColorOrder::BGR);

if(resolution == "720p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_720_P;
} else if(resolution == "400p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_400_P;
} else if(resolution == "800p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_800_P;
} else if(resolution == "480p") {
monoResolution = dai::node::MonoCamera::Properties::SensorResolution::THE_480_P;
} else {
RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "Invalid parameter. -> monoResolution: %s", resolution.c_str());
throw std::runtime_error("Invalid mono camera resolution.");
}

monoLeft->setResolution(monoResolution);
monoLeft->setBoardSocket(dai::CameraBoardSocket::LEFT);
monoRight->setResolution(monoResolution);
monoRight->setBoardSocket(dai::CameraBoardSocket::RIGHT);

/// setting node configs
stereo->initialConfig.setConfidenceThreshold(confidence);
stereo->setRectifyEdgeFillColor(0); // black, to better see the cutout
stereo->initialConfig.setLeftRightCheckThreshold(LRchecktresh);
stereo->setSubpixel(subpixel);
stereo->setDepthAlign(dai::CameraBoardSocket::RGB);

spatialDetectionNetwork->setBlobPath(nnPath);
spatialDetectionNetwork->setConfidenceThreshold(0.5f);
spatialDetectionNetwork->input.setBlocking(false);
spatialDetectionNetwork->setBoundingBoxScaleFactor(0.5);
spatialDetectionNetwork->setDepthLowerThreshold(100);
spatialDetectionNetwork->setDepthUpperThreshold(5000);

// yolo specific parameters
spatialDetectionNetwork->setNumClasses(80);
spatialDetectionNetwork->setCoordinateSize(4);
spatialDetectionNetwork->setAnchors({10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319});
spatialDetectionNetwork->setAnchorMasks({{"side13", {3, 4, 5}}, {"side26", {1, 2, 3}}});
spatialDetectionNetwork->setIouThreshold(0.5f);

// Link plugins CAM -> STEREO -> XLINK
monoLeft->out.link(stereo->left);
monoRight->out.link(stereo->right);

// Link plugins CAM -> NN -> XLINK
colorCam->preview.link(spatialDetectionNetwork->input);
if(syncNN)
spatialDetectionNetwork->passthrough.link(xoutRgb->input);
else
colorCam->preview.link(xoutRgb->input);

spatialDetectionNetwork->out.link(xoutNN->input);

stereo->depth.link(spatialDetectionNetwork->inputDepth);
spatialDetectionNetwork->passthroughDepth.link(xoutDepth->input);
}
else {
xoutRgb->setStreamName("rgb");
nnOut->setStreamName("detections");

// Properties
camRgb->setPreviewSize(416, 416);
camRgb->setResolution(dai::ColorCameraProperties::SensorResolution::THE_1080_P);
camRgb->setInterleaved(false);
camRgb->setColorOrder(dai::ColorCameraProperties::ColorOrder::BGR);
camRgb->setFps(40);

// Network specific settings
detectionNetwork->setConfidenceThreshold(0.5f);
detectionNetwork->setNumClasses(80);
detectionNetwork->setCoordinateSize(4);
detectionNetwork->setAnchors({10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319});
detectionNetwork->setAnchorMasks({{"side26", {1, 2, 3}}, {"side13", {3, 4, 5}}});
detectionNetwork->setIouThreshold(0.5f);
detectionNetwork->setBlobPath(nnPath);
detectionNetwork->setNumInferenceThreads(2);
detectionNetwork->input.setBlocking(false);

// Linking
camRgb->preview.link(detectionNetwork->input);
if(syncNN)
detectionNetwork->passthrough.link(xoutRgb->input);
else
camRgb->preview.link(xoutRgb->input);

detectionNetwork->out.link(nnOut->input);
}

monoLeft->setResolution(monoResolution);
monoLeft->setBoardSocket(dai::CameraBoardSocket::LEFT);
monoRight->setResolution(monoResolution);
monoRight->setBoardSocket(dai::CameraBoardSocket::RIGHT);

/// setting node configs
stereo->initialConfig.setConfidenceThreshold(confidence);
stereo->setRectifyEdgeFillColor(0); // black, to better see the cutout
stereo->initialConfig.setLeftRightCheckThreshold(LRchecktresh);
stereo->setSubpixel(subpixel);
stereo->setDepthAlign(dai::CameraBoardSocket::RGB);

spatialDetectionNetwork->setBlobPath(nnPath);
spatialDetectionNetwork->setConfidenceThreshold(0.5f);
spatialDetectionNetwork->input.setBlocking(false);
spatialDetectionNetwork->setBoundingBoxScaleFactor(0.5);
spatialDetectionNetwork->setDepthLowerThreshold(100);
spatialDetectionNetwork->setDepthUpperThreshold(5000);

// yolo specific parameters
spatialDetectionNetwork->setNumClasses(80);
spatialDetectionNetwork->setCoordinateSize(4);
spatialDetectionNetwork->setAnchors({10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319});
spatialDetectionNetwork->setAnchorMasks({{"side13", {3, 4, 5}}, {"side26", {1, 2, 3}}});
spatialDetectionNetwork->setIouThreshold(0.5f);

// Link plugins CAM -> STEREO -> XLINK
monoLeft->out.link(stereo->left);
monoRight->out.link(stereo->right);

// Link plugins CAM -> NN -> XLINK
colorCam->preview.link(spatialDetectionNetwork->input);
if(syncNN)
spatialDetectionNetwork->passthrough.link(xoutRgb->input);
else
colorCam->preview.link(xoutRgb->input);

spatialDetectionNetwork->out.link(xoutNN->input);

stereo->depth.link(spatialDetectionNetwork->inputDepth);
spatialDetectionNetwork->passthroughDepth.link(xoutDepth->input);
return pipeline;
}

Expand All @@ -119,7 +159,7 @@ int main(int argc, char** argv) {
std::string tfPrefix, resourceBaseFolder, nnPath;
std::string camera_param_uri;
std::string nnName(BLOB_NAME); // Set your blob name for the model here
bool syncNN, subpixel;
bool syncNN, subpixel, spatial_camera;
int confidence = 200, LRchecktresh = 5;
std::string monoResolution = "400p";

Expand All @@ -141,6 +181,7 @@ int main(int argc, char** argv) {
node->get_parameter("LRchecktresh", LRchecktresh);
node->get_parameter("monoResolution", monoResolution);
node->get_parameter("resourceBaseFolder", resourceBaseFolder);
node->get_parameter("spatial_camera", spatial_camera);

if(resourceBaseFolder.empty()) {
throw std::runtime_error("Send the path to the resouce folder containing NNBlob in \'resourceBaseFolder\' ");
Expand All @@ -153,7 +194,7 @@ int main(int argc, char** argv) {
}

nnPath = resourceBaseFolder + "/" + nnName;
dai::Pipeline pipeline = createPipeline(syncNN, subpixel, nnPath, confidence, LRchecktresh, monoResolution);
dai::Pipeline pipeline = createPipeline(spatial_camera, syncNN, subpixel, nnPath, confidence, LRchecktresh, monoResolution);
dai::Device device(pipeline);

auto colorQueue = device.getOutputQueue("preview", 30, false);
Expand Down

0 comments on commit 29c2906

Please sign in to comment.