* Optimize the YOLOv26 depth interface.

This commit is contained in:
lxowalle
2026-08-04 17:38:52 +08:00
parent 82a2aecd13
commit cdb27db679
4 changed files with 384 additions and 83 deletions

8
.gitignore vendored
View File

@@ -82,4 +82,10 @@ projects/*/build/
examples/*/dist/ examples/*/dist/
projects/*/dist/ projects/*/dist/
examples/*/data/ examples/*/data/
projects/*/data/ projects/*/data/
.agents
.codex
.trellis
.gitattributes
AGENTS.md

View File

@@ -11,6 +11,8 @@
#include <array> #include <array>
#include <cmath> #include <cmath>
#include <cstdint> #include <cstdint>
#include <limits>
#include <memory>
#include <vector> #include <vector>
#include "maix_basic.hpp" #include "maix_basic.hpp"
#include "maix_nn.hpp" #include "maix_nn.hpp"
@@ -25,35 +27,39 @@ namespace maix::nn
* Monocular depth estimation model (yolo26n-depth etc.) inference wrapper. * Monocular depth estimation model (yolo26n-depth etc.) inference wrapper.
* *
* Usage: * Usage:
* nn::Yolo26Depth model("/tmp/yolo26n-depth.mud"); * nn::YOLO26Depth model("/tmp/yolo26n-depth.mud");
* image::Image *img = cam.read(); * image::Image *img = cam.read();
* image::Image *heatmap = model.get_depth_image(*img, image::FIT_CONTAIN, image::CMap::JET); * tensor::Tensor *depth = model.get_depth(*img, image::FIT_CONTAIN);
* image::Image *heatmap = model.depth_to_image(*depth, img->width(), img->height());
* float meters = model.get_distance(*depth, img->width() / 2, img->height() / 2,
* img->width(), img->height());
* disp.show(*heatmap); * disp.show(*heatmap);
* delete heatmap; * delete heatmap;
* delete depth;
* delete img;
* *
* Note: the model MUD file's [extra] section should have model_type=yolo26_depth, input_type=rgb. * Note: the model MUD file's [extra] section should have model_type=yolo26_depth, input_type=rgb.
* @maixpy maix.nn.Yolo26Depth * @maixpy maix.nn.YOLO26Depth
*/ */
class Yolo26Depth class YOLO26Depth
{ {
public: public:
/** /**
* Construct a new Yolo26Depth object * Construct a new YOLO26Depth object
* @param model MUD model path, if empty, will not load model, you can call load() later. * @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. * 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. * @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. * If you want to ensure every time forward output the input's result, set this arg to false please.
* Default true to ensure speed. * Default true to ensure speed.
* @maixpy maix.nn.Yolo26Depth.__init__ * @maixpy maix.nn.YOLO26Depth.__init__
*/ */
Yolo26Depth(const string &model = "", bool dual_buff = true) YOLO26Depth(const string &model = "", bool dual_buff = true)
{ {
_model = nullptr; _model = nullptr;
_dual_buff = dual_buff; _dual_buff = dual_buff;
_input_img_fmt = image::Format::FMT_RGB888; _input_img_fmt = image::Format::FMT_RGB888;
_input_w = _input_h = 0; _input_w = _input_h = 0;
_output_w = _output_h = 0; _output_w = _output_h = 0;
_cmap = image::CMap::JET;
if (!model.empty()) if (!model.empty())
{ {
err::Err e = load(model); err::Err e = load(model);
@@ -64,7 +70,7 @@ namespace maix::nn
} }
} }
~Yolo26Depth() ~YOLO26Depth()
{ {
if (_model) if (_model)
{ {
@@ -80,7 +86,7 @@ namespace maix::nn
* - input_type: rgb or bgr * - input_type: rgb or bgr
* @param model MUD model path * @param model MUD model path
* @return error code, if load failed, return error code * @return error code, if load failed, return error code
* @maixpy maix.nn.Yolo26Depth.load * @maixpy maix.nn.YOLO26Depth.load
*/ */
err::Err load(const string &model) err::Err load(const string &model)
{ {
@@ -128,7 +134,7 @@ namespace maix::nn
/** /**
* Get model input size, only for image input * Get model input size, only for image input
* @return model input size * @return model input size
* @maixpy maix.nn.Yolo26Depth.input_size * @maixpy maix.nn.YOLO26Depth.input_size
*/ */
image::Size input_size() image::Size input_size()
{ {
@@ -138,7 +144,7 @@ namespace maix::nn
/** /**
* Get model input width, only for image input * Get model input width, only for image input
* @return model input size of width * @return model input size of width
* @maixpy maix.nn.Yolo26Depth.input_width * @maixpy maix.nn.YOLO26Depth.input_width
*/ */
int input_width() int input_width()
{ {
@@ -148,7 +154,7 @@ namespace maix::nn
/** /**
* Get model input height, only for image input * Get model input height, only for image input
* @return model input size of height * @return model input size of height
* @maixpy maix.nn.Yolo26Depth.input_height * @maixpy maix.nn.YOLO26Depth.input_height
*/ */
int input_height() int input_height()
{ {
@@ -158,7 +164,7 @@ namespace maix::nn
/** /**
* Get input image format, only for image input * Get input image format, only for image input
* @return input image format, image::Format type. * @return input image format, image::Format type.
* @maixpy maix.nn.Yolo26Depth.input_format * @maixpy maix.nn.YOLO26Depth.input_format
*/ */
image::Format input_format() image::Format input_format()
{ {
@@ -168,7 +174,7 @@ namespace maix::nn
/** /**
* Get model output size (depth map size), only for image input * Get model output size (depth map size), only for image input
* @return model output size * @return model output size
* @maixpy maix.nn.Yolo26Depth.output_size * @maixpy maix.nn.YOLO26Depth.output_size
*/ */
image::Size output_size() image::Size output_size()
{ {
@@ -183,23 +189,22 @@ namespace maix::nn
* Default Fit.FIT_CONTAIN, see image.Fit. * 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. * @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. * @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 * @maixpy maix.nn.YOLO26Depth.get_depth
*/ */
tensor::Tensor *get_depth(image::Image &img, image::Fit fit = image::FIT_CONTAIN) tensor::Tensor *get_depth(image::Image &img, image::Fit fit = image::FIT_CONTAIN)
{ {
if (_model == nullptr) return nullptr; if (_model == nullptr) return nullptr;
tensor::Tensors *outputs = _model->forward_image(img, {}, {}, fit, false, false); std::unique_ptr<tensor::Tensors> outputs(_model->forward_image(img, {}, {}, fit, false, false));
if (!outputs) return nullptr; if (!outputs) return nullptr;
if (outputs->size() == 0)
throw err::Exception(err::ERR_RUNTIME, "depth model returned no output tensor");
tensor::Tensor *t = outputs->begin()->second; tensor::Tensor *t = outputs->begin()->second;
if (t->dtype() != tensor::DType::FLOAT32) if (!t)
{ throw err::Exception(err::ERR_RUNTIME, "depth model returned a null output tensor");
delete outputs; DepthInfo info = _get_depth_info(*t);
return nullptr; _output_h = info.height;
} _output_w = info.width;
_output_h = t->shape()[2];
_output_w = t->shape()[3];
tensor::Tensor *result = new tensor::Tensor(t->shape(), t->dtype(), t->data(), true); tensor::Tensor *result = new tensor::Tensor(t->shape(), t->dtype(), t->data(), true);
delete outputs;
return result; return result;
} }
@@ -214,86 +219,192 @@ namespace maix::nn
* Default image.CMap.JET (near red/yellow, far blue). * 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. * @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. * @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 * @maixpy maix.nn.YOLO26Depth.get_depth_image
*/ */
image::Image *get_depth_image(image::Image &img, image::Fit fit = image::FIT_CONTAIN, image::Image *get_depth_image(image::Image &img, image::Fit fit = image::FIT_CONTAIN,
image::CMap cmap = image::CMap::JET) image::CMap cmap = image::CMap::JET)
{ {
if (_model == nullptr) return nullptr; std::unique_ptr<tensor::Tensor> depth(get_depth(img, fit));
_cmap = cmap; if (!depth) return nullptr;
return depth_to_image(*depth, img.width(), img.height(), fit, cmap);
}
tensor::Tensors *outputs = _model->forward_image(img, {}, {}, fit, false, false); /**
if (!outputs) return nullptr; * Convert an existing raw depth tensor to an RGB heatmap without running inference.
tensor::Tensor *t = outputs->begin()->second; * @param depth raw float32 depth tensor returned by get_depth().
if (t->dtype() != tensor::DType::FLOAT32) * @param image_width width of the source image passed to get_depth().
* @param image_height height of the source image passed to get_depth().
* @param fit resize fit mode used by get_depth().
* @param cmap color map used to visualize the depth values.
* @return a newly allocated image.Image at source-image size. In C++, delete it after use.
* @throw If the tensor, image size, or fit mode is invalid, raises err.Exception.
* @maixpy maix.nn.YOLO26Depth.depth_to_image
* @maixcdk maix.nn.YOLO26Depth.depth_to_image
*/
image::Image *depth_to_image(tensor::Tensor &depth, int image_width, int image_height,
image::Fit fit = image::FIT_CONTAIN,
image::CMap cmap = image::CMap::JET)
{
DepthInfo info = _get_depth_info(depth);
DepthTransform transform = _get_depth_transform(image_width, image_height,
info.width, info.height, fit);
const float *color_data = info.data;
int color_width = info.width;
int color_height = info.height;
std::vector<float> content;
if (fit == image::FIT_CONTAIN &&
(transform.content_width != info.width || transform.content_height != info.height))
{ {
delete outputs; color_width = transform.content_width;
return nullptr; color_height = transform.content_height;
content.resize(color_width * color_height);
for (int y = 0; y < color_height; ++y)
{
const float *src = info.data + (y + transform.offset_y) * info.width + transform.offset_x;
std::copy(src, src + color_width, content.begin() + y * color_width);
}
color_data = content.data();
} }
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 image::Image *heatmap = _colorize(color_data, color_width, color_height, cmap);
int top = 0, bottom = 0, left = 0, right = 0; if (color_width == image_width && color_height == image_height)
if (fit == image::Fit::FIT_CONTAIN && (img.width() != out_w || img.height() != out_h)) return heatmap;
{
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) image::Fit restore_fit = (fit == image::FIT_COVER) ? image::FIT_CONTAIN : image::FIT_FILL;
std::vector<float> crop(crop_w * crop_h); image::Image *result = heatmap->resize(image_width, image_height, restore_fit);
for (int y = 0; y < crop_h; y++) delete heatmap;
for (int x = 0; x < crop_w; x++) return result;
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); * Read the estimated distance at a source-image position from an existing depth tensor.
* This method only performs CPU-side coordinate mapping and does not run inference.
* @param depth raw float32 depth tensor returned by get_depth().
* @param x x coordinate in the source image.
* @param y y coordinate in the source image.
* @param image_width width of the source image passed to get_depth().
* @param image_height height of the source image passed to get_depth().
* @param fit resize fit mode used by get_depth().
* @return estimated distance in meters, or NaN if the coordinate is outside the
* represented image region or the model value is invalid.
* @throw If the tensor, image size, or fit mode is invalid, raises err.Exception.
* @maixpy maix.nn.YOLO26Depth.get_distance
* @maixcdk maix.nn.YOLO26Depth.get_distance
*/
float get_distance(tensor::Tensor &depth, int x, int y, int image_width, int image_height,
image::Fit fit = image::FIT_CONTAIN)
{
DepthInfo info = _get_depth_info(depth);
DepthTransform transform = _get_depth_transform(image_width, image_height,
info.width, info.height, fit);
if (x < 0 || y < 0 || x >= image_width || y >= image_height)
return std::numeric_limits<float>::quiet_NaN();
// resize back to input image size int depth_x = transform.offset_x + (int)(x * transform.scale_x);
if (crop_w != img.width() || crop_h != img.height()) int depth_y = transform.offset_y + (int)(y * transform.scale_y);
{ if (depth_x < 0 || depth_y < 0 || depth_x >= info.width || depth_y >= info.height)
image::Image *result = heatmap->resize(img.width(), img.height(), image::FIT_FILL); return std::numeric_limits<float>::quiet_NaN();
delete heatmap;
delete outputs; float value = info.data[depth_y * info.width + depth_x];
return result; if (!std::isfinite(value) || value <= 0)
} return std::numeric_limits<float>::quiet_NaN();
delete outputs; return value;
return heatmap;
} }
private: private:
struct DepthInfo
{
const float *data;
int width;
int height;
};
struct DepthTransform
{
float scale_x;
float scale_y;
int offset_x;
int offset_y;
int content_width;
int content_height;
};
nn::NN *_model; nn::NN *_model;
bool _dual_buff; bool _dual_buff;
image::Format _input_img_fmt; image::Format _input_img_fmt;
int _input_w, _input_h; int _input_w, _input_h;
int _output_w, _output_h; int _output_w, _output_h;
image::CMap _cmap;
DepthInfo _get_depth_info(tensor::Tensor &depth)
{
if (depth.dtype() != tensor::DType::FLOAT32)
throw err::Exception(err::ERR_ARGS, "depth tensor dtype must be float32");
std::vector<int> shape = depth.shape();
if (shape.size() != 4 || shape[0] != 1 || shape[1] != 1 ||
shape[2] <= 0 || shape[3] <= 0 || depth.data() == nullptr)
throw err::Exception(err::ERR_ARGS, "depth tensor shape must be [1, 1, H, W]");
return {(const float *)depth.data(), shape[3], shape[2]};
}
DepthTransform _get_depth_transform(int image_width, int image_height,
int depth_width, int depth_height, image::Fit fit)
{
if (image_width <= 0 || image_height <= 0)
throw err::Exception(err::ERR_ARGS, "source image size must be positive");
DepthTransform transform = {};
if (fit == image::FIT_FILL)
{
transform.content_width = depth_width;
transform.content_height = depth_height;
transform.offset_x = 0;
transform.offset_y = 0;
}
else if (fit == image::FIT_CONTAIN)
{
float scale = std::min((float)depth_width / image_width,
(float)depth_height / image_height);
transform.content_width = (int)std::round(image_width * scale);
transform.content_height = (int)std::round(image_height * scale);
transform.offset_x = (depth_width - transform.content_width) / 2;
transform.offset_y = (depth_height - transform.content_height) / 2;
}
else if (fit == image::FIT_COVER)
{
float scale = std::max((float)depth_width / image_width,
(float)depth_height / image_height);
transform.content_width = (int)std::round(image_width * scale);
transform.content_height = (int)std::round(image_height * scale);
transform.offset_x = -(transform.content_width - depth_width) / 2;
transform.offset_y = -(transform.content_height - depth_height) / 2;
}
else
{
throw err::Exception(err::ERR_ARGS, "unsupported depth image fit mode");
}
if (transform.content_width <= 0 || transform.content_height <= 0)
throw err::Exception(err::ERR_ARGS, "invalid mapped depth content size");
transform.scale_x = (float)transform.content_width / image_width;
transform.scale_y = (float)transform.content_height / image_height;
return transform;
}
/** /**
* Convert depth map to RGB888 heatmap: disparity(1/depth) + min-max normalize + cmap lookup. * Convert depth map to RGB888 heatmap: disparity(1/depth) + min-max normalize + cmap lookup.
* Return a new image::Image, caller should delete it. * Return a new image::Image, caller should delete it.
*/ */
image::Image *_colorize(const float *depth, int w, int h) image::Image *_colorize(const float *depth, int w, int h, image::CMap cmap)
{ {
auto &colors = image::cmap_colors_rgb(_cmap); auto &colors = image::cmap_colors_rgb(cmap);
image::Image *result = new image::Image(w, h, image::Format::FMT_RGB888); image::Image *result = new image::Image(w, h, image::Format::FMT_RGB888);
uint8_t *img_data = (uint8_t *)result->data(); uint8_t *img_data = (uint8_t *)result->data();
int n = w * h; int n = w * h;
float min_v = FLT_MAX, max_v = -FLT_MAX; float min_v = FLT_MAX, max_v = -FLT_MAX;
int valid_count = 0;
for (int i = 0; i < n; i++) for (int i = 0; i < n; i++)
{ {
float d = depth[i]; float d = depth[i];
@@ -302,8 +413,14 @@ namespace maix::nn
float v = 1.0f / d; // disparity float v = 1.0f / d; // disparity
if (v < min_v) min_v = v; if (v < min_v) min_v = v;
if (v > max_v) max_v = v; if (v > max_v) max_v = v;
++valid_count;
} }
} }
if (valid_count == 0)
{
memset(img_data, 0, n * 3);
return result;
}
if (min_v == max_v) if (min_v == max_v)
{ {
memset(img_data, 127, n * 3); memset(img_data, 127, n * 3);

View File

@@ -18,12 +18,38 @@ Single image inference, save heatmap:
./nn_yolo26_depth /tmp/yolo26n-depth.mud /tmp/bus.jpg /tmp/bus_heatmap.jpg ./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): Real-time camera inference with an interactive depth heatmap (exit with SIGINT):
```bash ```bash
./nn_yolo26_depth /tmp/yolo26n-depth.mud ./nn_yolo26_depth /tmp/yolo26n-depth.mud
``` ```
### Touch distance probes
- Tap inside the displayed heatmap to add a distance probe.
- A probe stays at the selected image position and its value updates on every new frame.
- Up to five probes are kept. Adding a sixth removes the oldest probe.
- Tap within 20 screen pixels of an existing probe to remove it.
- Taps in black letterbox padding are ignored.
- Invalid model samples are displayed as `N/A`.
Distances are monocular depth estimates in meters. The released YOLO26-depth weights include global metric-scale calibration, but the values are not a replacement for a calibrated distance sensor.
The example runs one NPU inference per processed frame. Heatmap generation and all active probe values reuse that same depth tensor on the CPU, so adding probes does not add NPU inference calls.
The reusable C++ flow is:
```cpp
nn::YOLO26Depth model("/tmp/yolo26n-depth.mud");
tensor::Tensor *depth = model.get_depth(frame, image::FIT_CONTAIN); // NPU once
image::Image *heatmap = model.depth_to_image(
*depth, frame.width(), frame.height(), image::FIT_CONTAIN, image::CMap::JET);
float meters = model.get_distance(
*depth, x, y, frame.width(), frame.height(), image::FIT_CONTAIN); // CPU only
delete heatmap;
delete depth;
```
## Model files ## Model files
Put the model on the device first: Put the model on the device first:

View File

@@ -7,6 +7,7 @@
* 例: ./nn_yolo26_depth /tmp/yolo26n-depth.mud /tmp/bus.jpg /tmp/bus_heatmap.jpg * 例: ./nn_yolo26_depth /tmp/yolo26n-depth.mud /tmp/bus.jpg /tmp/bus_heatmap.jpg
* - 实时摄像头推理并显示热力图: * - 实时摄像头推理并显示热力图:
* ./nn_yolo26_depth mud_model_path * ./nn_yolo26_depth mud_model_path
* 点击画面可添加/取消距离探针,最多保留 5 个,距离随画面实时更新。
* *
* 依赖(设备上需存在): * 依赖(设备上需存在):
* /tmp/yolo26n-depth.mud * /tmp/yolo26n-depth.mud
@@ -21,10 +22,135 @@
#include "maix_image.hpp" #include "maix_image.hpp"
#include "maix_nn_yolo26_depth.hpp" #include "maix_nn_yolo26_depth.hpp"
#include "maix_time.hpp" #include "maix_time.hpp"
#include "maix_touchscreen.hpp"
#include "main.h" #include "main.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <memory>
#include <vector>
using namespace maix; using namespace maix;
namespace
{
constexpr size_t MAX_PROBES = 5;
constexpr int PROBE_HIT_RADIUS = 20;
constexpr int PROBE_HIT_RADIUS_SQUARED = PROBE_HIT_RADIUS * PROBE_HIT_RADIUS;
struct Probe
{
int x;
int y;
};
struct DisplayTransform
{
float scale;
int offset_x;
int offset_y;
int content_width;
int content_height;
};
DisplayTransform get_display_transform(int image_width, int image_height,
int display_width, int display_height)
{
float scale = std::min((float)display_width / image_width,
(float)display_height / image_height);
int content_width = (int)std::round(image_width * scale);
int content_height = (int)std::round(image_height * scale);
return {scale,
(display_width - content_width) / 2,
(display_height - content_height) / 2,
content_width,
content_height};
}
bool screen_to_image(const DisplayTransform &transform, int screen_x, int screen_y,
int image_width, int image_height, Probe &point)
{
if (screen_x < transform.offset_x || screen_y < transform.offset_y ||
screen_x >= transform.offset_x + transform.content_width ||
screen_y >= transform.offset_y + transform.content_height)
return false;
point.x = std::clamp((int)((screen_x - transform.offset_x) / transform.scale),
0, image_width - 1);
point.y = std::clamp((int)((screen_y - transform.offset_y) / transform.scale),
0, image_height - 1);
return true;
}
Probe image_to_screen(const DisplayTransform &transform, const Probe &point)
{
return {transform.offset_x + (int)(point.x * transform.scale),
transform.offset_y + (int)(point.y * transform.scale)};
}
void update_probes(std::vector<Probe> &probes, int screen_x, int screen_y,
int image_width, int image_height,
int display_width, int display_height)
{
DisplayTransform transform = get_display_transform(image_width, image_height,
display_width, display_height);
Probe new_probe;
if (!screen_to_image(transform, screen_x, screen_y, image_width, image_height, new_probe))
return;
size_t nearest = probes.size();
int nearest_distance = PROBE_HIT_RADIUS_SQUARED + 1;
for (size_t i = 0; i < probes.size(); ++i)
{
Probe screen_probe = image_to_screen(transform, probes[i]);
int dx = screen_x - screen_probe.x;
int dy = screen_y - screen_probe.y;
int distance = dx * dx + dy * dy;
if (distance <= PROBE_HIT_RADIUS_SQUARED && distance < nearest_distance)
{
nearest = i;
nearest_distance = distance;
}
}
if (nearest < probes.size())
{
probes.erase(probes.begin() + nearest);
return;
}
if (probes.size() == MAX_PROBES)
probes.erase(probes.begin());
probes.push_back(new_probe);
}
void draw_probe(image::Image &heatmap, const Probe &probe, float distance)
{
const image::Color marker_color = image::Color::from_rgb(255, 255, 255);
heatmap.draw_circle(probe.x, probe.y, 8, image::COLOR_BLACK, -1);
heatmap.draw_circle(probe.x, probe.y, 5, marker_color, -1);
char label[32];
if (std::isfinite(distance))
std::snprintf(label, sizeof(label), "%.2f m", distance);
else
std::snprintf(label, sizeof(label), "N/A");
image::Size text_size = image::string_size(label);
int label_x = probe.x + 10;
int label_y = probe.y - text_size.height() - 8;
if (label_y < 2)
label_y = probe.y + 10;
label_x = std::clamp(label_x, 2, std::max(2, heatmap.width() - text_size.width() - 2));
label_y = std::clamp(label_y, 2, std::max(2, heatmap.height() - text_size.height() - 2));
heatmap.draw_rect(label_x - 2, label_y - 2,
text_size.width() + 4, text_size.height() + 4,
image::COLOR_BLACK, -1);
heatmap.draw_string(label_x, label_y, label, marker_color);
}
}
int _main(int argc, char *argv[]) int _main(int argc, char *argv[])
{ {
int ret = 0; int ret = 0;
@@ -42,7 +168,7 @@ int _main(int argc, char *argv[])
image::CMap cmap = image::CMap::JET; image::CMap cmap = image::CMap::JET;
log::info("model path: %s", model_path); log::info("model path: %s", model_path);
nn::Yolo26Depth model(model_path, dual_buff); nn::YOLO26Depth model(model_path, dual_buff);
log::info("load model %s success", model_path); log::info("load model %s success", model_path);
log::info("model input size: %dx%d, format: %d", log::info("model input size: %dx%d, format: %d",
model.input_width(), model.input_height(), (int)model.input_format()); model.input_width(), model.input_height(), (int)model.input_format());
@@ -68,7 +194,7 @@ int _main(int argc, char *argv[])
err::check_null_raise(heatmap, "get_depth_image failed"); err::check_null_raise(heatmap, "get_depth_image failed");
// 保存热力图 // 保存热力图
err::Err e = heatmap->save(output_path); err::Err e = heatmap->save(output_path.c_str());
if (e != err::ERR_NONE) if (e != err::ERR_NONE)
{ {
log::error("save heatmap to %s failed", output_path.c_str()); log::error("save heatmap to %s failed", output_path.c_str());
@@ -87,29 +213,55 @@ int _main(int argc, char *argv[])
// ---------- 模式二: 实时摄像头推理并显示热力图 ---------- // ---------- 模式二: 实时摄像头推理并显示热力图 ----------
camera::Camera cam = camera::Camera(model.input_width(), model.input_height(), model.input_format()); camera::Camera cam = camera::Camera(model.input_width(), model.input_height(), model.input_format());
display::Display disp = display::Display(); display::Display disp = display::Display();
touchscreen::TouchScreen touch = touchscreen::TouchScreen();
log::info("camera %dx%d", cam.width(), cam.height()); log::info("camera %dx%d", cam.width(), cam.height());
uint64_t frame_id = 0; uint64_t frame_id = 0;
uint64_t t_last = time::ticks_ms(); uint64_t t_last = time::ticks_ms();
int fps = 0; int fps = 0;
bool was_pressed = false;
std::vector<Probe> probes;
while (!app::need_exit()) while (!app::need_exit())
{ {
uint64_t t0 = time::ticks_ms(); uint64_t t0 = time::ticks_ms();
image::Image *img = cam.read(); std::unique_ptr<image::Image> img(cam.read());
if (!img) if (!img)
{ {
time::sleep_ms(10); time::sleep_ms(10);
continue; continue;
} }
image::Image *heatmap = model.get_depth_image(*img, image::FIT_CONTAIN, cmap); int touch_x = 0;
int touch_y = 0;
bool pressed = false;
err::Err touch_err = touch.read(touch_x, touch_y, pressed);
if (touch_err == err::ERR_NONE)
{
if (was_pressed && !pressed)
{
update_probes(probes, touch_x, touch_y,
img->width(), img->height(), disp.width(), disp.height());
}
was_pressed = pressed;
}
std::unique_ptr<tensor::Tensor> depth(model.get_depth(*img, image::FIT_CONTAIN));
if (!depth)
continue;
std::unique_ptr<image::Image> heatmap(
model.depth_to_image(*depth, img->width(), img->height(), image::FIT_CONTAIN, cmap));
if (heatmap) if (heatmap)
{ {
for (const Probe &probe : probes)
{
float distance = model.get_distance(*depth, probe.x, probe.y,
img->width(), img->height(), image::FIT_CONTAIN);
draw_probe(*heatmap, probe, distance);
}
disp.show(*heatmap); disp.show(*heatmap);
delete heatmap;
} }
delete img;
frame_id++; frame_id++;
if (frame_id % 20 == 0) if (frame_id % 20 == 0)