diff --git a/components/nn/include/maix_nn_yolo26_depth.hpp b/components/nn/include/maix_nn_yolo26_depth.hpp new file mode 100644 index 000000000..ac95eb438 --- /dev/null +++ b/components/nn/include/maix_nn_yolo26_depth.hpp @@ -0,0 +1,336 @@ +/** + * @author maixpy_skill + * @copyright Sipeed Ltd 2023- + * @license Apache 2.0 + * @update 2026.8.2: Add yolo26-depth support + */ + +#pragma once +#include +#include +#include +#include +#include +#include +#include "maix_basic.hpp" +#include "maix_nn.hpp" +#include "maix_image.hpp" +#include "maix_image_cmap.hpp" +#include "maix_time.hpp" + +namespace maix::nn +{ + /** + * YOLO26-depth + * Monocular depth estimation model (yolo26n-depth etc.) inference wrapper. + * + * Usage: + * nn::Yolo26Depth model("/tmp/yolo26n-depth.mud"); + * image::Image *img = cam.read(); + * image::Image *heatmap = model.get_depth_image(*img, image::FIT_CONTAIN, image::CMap::JET); + * disp.show(*heatmap); + * delete heatmap; + * + * Note: the model MUD file's [extra] section should have model_type=yolo26_depth, input_type=rgb. + * @maixpy maix.nn.Yolo26Depth + */ + class Yolo26Depth + { + public: + /** + * Construct a new Yolo26Depth object + * @param model MUD model path, if empty, will not load model, you can call load() later. + * if not empty, will load model and will raise err::Exception if load failed. + * @param[in] dual_buff prepare dual input output buffer to accelarate forward, that is, when NPU is forwarding we not wait and prepare the next input buff. + * If you want to ensure every time forward output the input's result, set this arg to false please. + * Default true to ensure speed. + * @maixpy maix.nn.Yolo26Depth.__init__ + */ + Yolo26Depth(const string &model = "", bool dual_buff = true) + { + _model = nullptr; + _dual_buff = dual_buff; + _input_img_fmt = image::Format::FMT_RGB888; + _input_w = _input_h = 0; + _output_w = _output_h = 0; + _cmap = image::CMap::JET; + if (!model.empty()) + { + err::Err e = load(model); + if (e != err::ERR_NONE) + { + throw err::Exception(e, "load model failed"); + } + } + } + + ~Yolo26Depth() + { + if (_model) + { + delete _model; + _model = nullptr; + } + } + + /** + * Load model from file, model format is .mud, + * MUD file should contain [extra] section, have key-values: + * - model_type: yolo26_depth + * - input_type: rgb or bgr + * @param model MUD model path + * @return error code, if load failed, return error code + * @maixpy maix.nn.Yolo26Depth.load + */ + err::Err load(const string &model) + { + if (_model) + { + delete _model; + _model = nullptr; + } + _model = new nn::NN(model, _dual_buff); + if (!_model) + { + return err::ERR_NO_MEM; + } + auto inputs = _model->inputs_info(); + if (inputs.empty()) + { + return err::ERR_ARGS; + } + // NHWC: [1, H, W, 3]; NCHW: [1, 3, H, W] + if (inputs[0].shape[3] <= 4) // NHWC + { + _input_h = inputs[0].shape[1]; + _input_w = inputs[0].shape[2]; + } + else + { + _input_h = inputs[0].shape[2]; + _input_w = inputs[0].shape[3]; + } + // input_type decides input image format + auto extra = _model->extra_info(); + auto it = extra.find("input_type"); + if (it != extra.end()) + { + if (it->second == "bgr") + _input_img_fmt = image::Format::FMT_BGR888; + else if (it->second == "gray") + _input_img_fmt = image::Format::FMT_GRAYSCALE; + else + _input_img_fmt = image::Format::FMT_RGB888; + } + return err::ERR_NONE; + } + + /** + * Get model input size, only for image input + * @return model input size + * @maixpy maix.nn.Yolo26Depth.input_size + */ + image::Size input_size() + { + return image::Size(_input_w, _input_h); + } + + /** + * Get model input width, only for image input + * @return model input size of width + * @maixpy maix.nn.Yolo26Depth.input_width + */ + int input_width() + { + return _input_w; + } + + /** + * Get model input height, only for image input + * @return model input size of height + * @maixpy maix.nn.Yolo26Depth.input_height + */ + int input_height() + { + return _input_h; + } + + /** + * Get input image format, only for image input + * @return input image format, image::Format type. + * @maixpy maix.nn.Yolo26Depth.input_format + */ + image::Format input_format() + { + return _input_img_fmt; + } + + /** + * Get model output size (depth map size), only for image input + * @return model output size + * @maixpy maix.nn.Yolo26Depth.output_size + */ + image::Size output_size() + { + return image::Size(_output_w, _output_h); + } + + /** + * Forward model and get raw image depth estimation data. + * @param img image, format should match model input_type, or will raise err.Exception + * @param fit image resize fit mode if input image not equal to model' input size, + * will auto resize to model's input size then detect, and recover to image input size. + * Default Fit.FIT_CONTAIN, see image.Fit. + * @throw If error occurred, will raise err::Exception, you can find reason in log, mostly caused by args error or hardware error. + * @return result, a tensor.Tensor object. If in dual_buff mode, value can be None(in Python) or nullptr(in C++) when not ready. In C++, you need to delete it after use. + * @maixpy maix.nn.Yolo26Depth.get_depth + */ + tensor::Tensor *get_depth(image::Image &img, image::Fit fit = image::FIT_CONTAIN) + { + if (_model == nullptr) return nullptr; + tensor::Tensors *outputs = _model->forward_image(img, {}, {}, fit, false, false); + if (!outputs) return nullptr; + tensor::Tensor *t = outputs->begin()->second; + if (t->dtype() != tensor::DType::FLOAT32) + { + delete outputs; + return nullptr; + } + _output_h = t->shape()[2]; + _output_w = t->shape()[3]; + tensor::Tensor *result = new tensor::Tensor(t->shape(), t->dtype(), t->data(), true); + delete outputs; + return result; + } + + /** + * Forward model and get image depth estimation data normlized to [0, 255] and as a image.Image object. + * @param img image, format should match model input_type, or will raise err.Exception + * @param fit image resize fit mode if input image not equal to model' input size, + * will auto resize to model's input size then detect, and recover to image input size. + * Default Fit.FIT_CONTAIN, see image.Fit. + * @param cmap Color map used convert grayscale distance estimation image to RGB image. + * Diiferent cmap will influence finally image. + * Default image.CMap.JET (near red/yellow, far blue). + * @throw If error occurred, will raise err::Exception, you can find reason in log, mostly caused by args error or hardware error. + * @return result, a image::Image object. If in dual_buff mode, value can be None(in Python) or nullptr(in C++) when not ready. In C++, you need to delete it after use. + * @maixpy maix.nn.Yolo26Depth.get_depth_image + */ + image::Image *get_depth_image(image::Image &img, image::Fit fit = image::FIT_CONTAIN, + image::CMap cmap = image::CMap::JET) + { + if (_model == nullptr) return nullptr; + _cmap = cmap; + + tensor::Tensors *outputs = _model->forward_image(img, {}, {}, fit, false, false); + if (!outputs) return nullptr; + tensor::Tensor *t = outputs->begin()->second; + if (t->dtype() != tensor::DType::FLOAT32) + { + delete outputs; + return nullptr; + } + int out_w = t->shape()[3]; + int out_h = t->shape()[2]; + _output_w = out_w; + _output_h = out_h; + const float *depth = (const float *)t->data(); + + // compute letterbox content region (remove padding), keep depth pixel-aligned with image + int top = 0, bottom = 0, left = 0, right = 0; + if (fit == image::Fit::FIT_CONTAIN && (img.width() != out_w || img.height() != out_h)) + { + float gain = std::min((float)out_h / img.height(), (float)out_w / img.width()); + int rw = (int)std::round(img.width() * gain); + int rh = (int)std::round(img.height() * gain); + int pw = out_w - rw; + int ph = out_h - rh; + left = (int)std::round(pw / 2.0f - 0.1f); + right = (int)std::round(pw / 2.0f + 0.1f); + top = (int)std::round(ph / 2.0f - 0.1f); + bottom = (int)std::round(ph / 2.0f + 0.1f); + } + int crop_w = out_w - left - right; + int crop_h = out_h - top - bottom; + + // extract content region depth (row-major) + std::vector crop(crop_w * crop_h); + for (int y = 0; y < crop_h; y++) + for (int x = 0; x < crop_w; x++) + crop[y * crop_w + x] = depth[(y + top) * out_w + (x + left)]; + + // generate heatmap on crop size + image::Image *heatmap = _colorize(crop.data(), crop_w, crop_h); + + // resize back to input image size + if (crop_w != img.width() || crop_h != img.height()) + { + image::Image *result = heatmap->resize(img.width(), img.height(), image::FIT_FILL); + delete heatmap; + delete outputs; + return result; + } + delete outputs; + return heatmap; + } + + private: + nn::NN *_model; + bool _dual_buff; + image::Format _input_img_fmt; + int _input_w, _input_h; + int _output_w, _output_h; + image::CMap _cmap; + + /** + * Convert depth map to RGB888 heatmap: disparity(1/depth) + min-max normalize + cmap lookup. + * Return a new image::Image, caller should delete it. + */ + image::Image *_colorize(const float *depth, int w, int h) + { + auto &colors = image::cmap_colors_rgb(_cmap); + image::Image *result = new image::Image(w, h, image::Format::FMT_RGB888); + uint8_t *img_data = (uint8_t *)result->data(); + int n = w * h; + + float min_v = FLT_MAX, max_v = -FLT_MAX; + for (int i = 0; i < n; i++) + { + float d = depth[i]; + if (std::isfinite(d) && d > 0) + { + float v = 1.0f / d; // disparity + if (v < min_v) min_v = v; + if (v > max_v) max_v = v; + } + } + if (min_v == max_v) + { + memset(img_data, 127, n * 3); + return result; + } + float scale = 255.0f / (max_v - min_v); + +#pragma omp parallel for + for (int i = 0; i < n; i++) + { + float d = depth[i]; + if (std::isfinite(d) && d > 0) + { + uint8_t gray = (uint8_t)std::clamp((1.0f / d - min_v) * scale, 0.0f, 255.0f); + const auto &rgb = colors[gray]; + img_data[3 * i + 0] = rgb[0]; + img_data[3 * i + 1] = rgb[1]; + img_data[3 * i + 2] = rgb[2]; + } + else + { + img_data[3 * i + 0] = 0; + img_data[3 * i + 1] = 0; + img_data[3 * i + 2] = 0; + } + } + return result; + } + }; +} diff --git a/examples/nn_yolo26_depth/.gitignore b/examples/nn_yolo26_depth/.gitignore new file mode 100644 index 000000000..777259b5c --- /dev/null +++ b/examples/nn_yolo26_depth/.gitignore @@ -0,0 +1,5 @@ +build +dist +.config.mk +.flash.conf.json +data diff --git a/examples/nn_yolo26_depth/README.md b/examples/nn_yolo26_depth/README.md new file mode 100644 index 000000000..3c378ed2d --- /dev/null +++ b/examples/nn_yolo26_depth/README.md @@ -0,0 +1,35 @@ +# YOLO26-depth example + +YOLO26-depth (monocular depth estimation) example for MaixCDK, based on the header-only class `maix_nn_yolo26_depth.hpp`. + +## Build + +```bash +cd examples/nn_yolo26_depth +export CMAKE_POLICY_VERSION_MINIMUM=3.5 # required on CMake >= 4.x, see FAQ +maixcdk build -p maixcam2 +``` + +## Usage + +Single image inference, save heatmap: + +```bash +./nn_yolo26_depth /tmp/yolo26n-depth.mud /tmp/bus.jpg /tmp/bus_heatmap.jpg +``` + +Real-time camera inference, display heatmap on screen (exit with SIGINT): + +```bash +./nn_yolo26_depth /tmp/yolo26n-depth.mud +``` + +## Model files + +Put the model on the device first: + +```bash +scp yolo26n-depth.mud yolo26n-depth_w8a8_mix.axmodel root@:/tmp/ +``` + +Build method please visit [MaixCDK](https://github.com/sipeed/MaixCDK). diff --git a/examples/nn_yolo26_depth/main/CMakeLists.txt b/examples/nn_yolo26_depth/main/CMakeLists.txt new file mode 100644 index 000000000..6619a51c7 --- /dev/null +++ b/examples/nn_yolo26_depth/main/CMakeLists.txt @@ -0,0 +1,16 @@ +############### Add include ################### +list(APPEND ADD_INCLUDE "include" + ) +list(APPEND ADD_PRIVATE_INCLUDE "") +############################################### + +############ Add source files ################# +append_srcs_dir(ADD_SRCS "src") # append source file in src dir to var ADD_SRCS +############################################### + +###### Add required/dependent components ###### +list(APPEND ADD_REQUIREMENTS basic nn vision) +############################################### + +# register component, DYNAMIC or SHARED flags will make component compiled to dynamic(shared) lib +register_component() diff --git a/examples/nn_yolo26_depth/main/Kconfig b/examples/nn_yolo26_depth/main/Kconfig new file mode 100644 index 000000000..0ee4d182c --- /dev/null +++ b/examples/nn_yolo26_depth/main/Kconfig @@ -0,0 +1,9 @@ +menu "Main Component Config" + config MAIN_COMPONENT_EACH_BUILD + bool "Each component build a separate shared library file" + default n + + config MAIN_COMPONENT_BUILD_TYPE + string "Main component build type" + default "Release" +endmenu diff --git a/examples/nn_yolo26_depth/main/include/main.h b/examples/nn_yolo26_depth/main/include/main.h new file mode 100644 index 000000000..ff6398284 --- /dev/null +++ b/examples/nn_yolo26_depth/main/include/main.h @@ -0,0 +1,3 @@ +#pragma once + +#define APP_VERSION "0.1.0" diff --git a/examples/nn_yolo26_depth/main/src/main.cpp b/examples/nn_yolo26_depth/main/src/main.cpp new file mode 100644 index 000000000..40f174c5b --- /dev/null +++ b/examples/nn_yolo26_depth/main/src/main.cpp @@ -0,0 +1,134 @@ +/** + * YOLO26-depth example + * + * Usage: + * - 单张图片推理并保存热力图: + * ./nn_yolo26_depth mud_model_path image_path [output_path] + * 例: ./nn_yolo26_depth /tmp/yolo26n-depth.mud /tmp/bus.jpg /tmp/bus_heatmap.jpg + * - 实时摄像头推理并显示热力图: + * ./nn_yolo26_depth mud_model_path + * + * 依赖(设备上需存在): + * /tmp/yolo26n-depth.mud + * /tmp/yolo26n-depth_w8a8_mix.axmodel + * + * 退出: 摄像头模式下发送 SIGINT/SIGTERM(kill -INT ) 或设备退出信号。 + */ + +#include "maix_basic.hpp" +#include "maix_camera.hpp" +#include "maix_display.hpp" +#include "maix_image.hpp" +#include "maix_nn_yolo26_depth.hpp" +#include "maix_time.hpp" +#include "main.h" + +using namespace maix; + +int _main(int argc, char *argv[]) +{ + int ret = 0; + log::info("Program start"); + std::string help = "Usage: " + std::string(argv[0]) + " mud_model_path [image_path [output_path]]"; + + if (argc < 2) + { + log::error(help.c_str()); + return -1; + } + + const char *model_path = argv[1]; + bool dual_buff = false; + image::CMap cmap = image::CMap::JET; + + log::info("model path: %s", model_path); + nn::Yolo26Depth model(model_path, dual_buff); + log::info("load model %s success", model_path); + log::info("model input size: %dx%d, format: %d", + model.input_width(), model.input_height(), (int)model.input_format()); + log::info("dual buff mode: %d", dual_buff); + + if (argc >= 3) + { + // ---------- 模式一: 单张图片推理, 保存热力图 ---------- + const char *img_path = argv[2]; + std::string output_path = (argc >= 4) ? argv[3] : "depth_heatmap.jpg"; + log::info("load image now"); + maix::image::Image *img = maix::image::load(img_path, model.input_format()); + err::check_null_raise(img, "load image " + std::string(img_path) + " failed"); + log::info("load image %s success: %s", img_path, img->to_str().c_str()); + if (img->width() != model.input_width() || img->height() != model.input_height()) + { + log::warn("image size not match model input size, will auto resize from %dx%d to %dx%d", + img->width(), img->height(), model.input_width(), model.input_height()); + } + + // 推理得到深度热力图(自动 resize 回原图尺寸) + maix::image::Image *heatmap = model.get_depth_image(*img, image::FIT_CONTAIN, cmap); + err::check_null_raise(heatmap, "get_depth_image failed"); + + // 保存热力图 + err::Err e = heatmap->save(output_path); + if (e != err::ERR_NONE) + { + log::error("save heatmap to %s failed", output_path.c_str()); + ret = -1; + } + else + { + log::info("save heatmap to %s success", output_path.c_str()); + } + + delete heatmap; + delete img; + } + else + { + // ---------- 模式二: 实时摄像头推理并显示热力图 ---------- + camera::Camera cam = camera::Camera(model.input_width(), model.input_height(), model.input_format()); + display::Display disp = display::Display(); + log::info("camera %dx%d", cam.width(), cam.height()); + + uint64_t frame_id = 0; + uint64_t t_last = time::ticks_ms(); + int fps = 0; + + while (!app::need_exit()) + { + uint64_t t0 = time::ticks_ms(); + image::Image *img = cam.read(); + if (!img) + { + time::sleep_ms(10); + continue; + } + + image::Image *heatmap = model.get_depth_image(*img, image::FIT_CONTAIN, cmap); + if (heatmap) + { + disp.show(*heatmap); + delete heatmap; + } + delete img; + + frame_id++; + if (frame_id % 20 == 0) + { + uint64_t now = time::ticks_ms(); + fps = (int)(20000.0f / (now - t_last)); + t_last = now; + log::info("frame=%llu e2e=%llu ms fps=%d", (unsigned long long)frame_id, + (unsigned long long)(time::ticks_ms() - t0), fps); + } + } + } + + log::info("Program exit"); + return ret; +} + +int main(int argc, char *argv[]) +{ + sys::register_default_signal_handle(); + CATCH_EXCEPTION_RUN_RETURN(_main, -1, argc, argv); +}