236 lines
13 KiB
C++
236 lines
13 KiB
C++
#include "bmp_loader.h"
|
||
#include "result_json.h"
|
||
#include "yolo_runtime.h"
|
||
#include "mask_dilation.h"
|
||
|
||
#include <iostream>
|
||
#include <memory>
|
||
#include <string_view>
|
||
#include <vector>
|
||
#include <limits>
|
||
|
||
static const char* value_after(int argc, char** argv, std::string_view option) {
|
||
for (int index = 1; index + 1 < argc; ++index)
|
||
if (argv[index] == option) return argv[index + 1];
|
||
return nullptr;
|
||
}
|
||
|
||
struct RuntimeDeleter {
|
||
void operator()(yolo_runtime_t* value) const { yolo_runtime_destroy(value); }
|
||
};
|
||
struct ResultDeleter {
|
||
void operator()(yolo_result_set_t* value) const { yolo_result_set_release(value); }
|
||
};
|
||
|
||
struct MaskBuffers {
|
||
std::vector<unsigned char*> data;
|
||
explicit MaskBuffers(size_t count):data(count,nullptr) {}
|
||
~MaskBuffers() {for(auto* p:data) free(p);}
|
||
};
|
||
|
||
int main(int argc, char** argv) {
|
||
int dilation_radius=0;
|
||
if(parse_mask_dilation(argc,argv,&dilation_radius)!=0) {std::cerr<<"Invalid --mask-dilation: integer 0..100 required\n";return 2;}
|
||
const char* config = value_after(argc, argv, "--config");
|
||
const char* image_path = value_after(argc, argv, "--image");
|
||
const char* json_path = value_after(argc, argv, "--json");
|
||
const char* geometry_name = value_after(argc, argv, "--seg-geometry");
|
||
if (!config || !image_path || !json_path) {
|
||
std::cerr << "Usage: yolo_cpp_example --config model.yaml --image sample.bmp "
|
||
"--json result.json [--seg-geometry rect|circle|mean|fit] [--rgb-mask] [--mask-dilation 0..100] "
|
||
"[--roi-lt L T | --roi-whlt W H L T]\n";
|
||
return 2;
|
||
}
|
||
if(geometry_name && std::string_view(geometry_name)!="rect" && std::string_view(geometry_name)!="circle" &&
|
||
std::string_view(geometry_name)!="mean" && std::string_view(geometry_name)!="fit") return 2;
|
||
if (yolo_runtime_get_abi_version() != YOLO_RUNTIME_ABI_VERSION) return 3;
|
||
unsigned char* raw_pixels = nullptr;
|
||
int width = 0, height = 0;
|
||
if (load_bmp_rgb(image_path, &raw_pixels, &width, &height) != 0) return 4;
|
||
std::unique_ptr<unsigned char, decltype(&free_bmp_rgb)> pixels(raw_pixels, free_bmp_rgb);
|
||
|
||
yolo_error_info_v1_t error{};
|
||
error.struct_size = sizeof(error);
|
||
error.version = YOLO_ERROR_INFO_VERSION;
|
||
yolo_runtime_t* raw_runtime = nullptr;
|
||
if (yolo_runtime_create(config, &raw_runtime, &error) != YOLO_STATUS_OK || !raw_runtime) {
|
||
std::cerr << "Create failed: " << error.message << '\n';
|
||
return 5;
|
||
}
|
||
std::unique_ptr<yolo_runtime_t, RuntimeDeleter> runtime(raw_runtime);
|
||
yolo_runtime_info_v1_t info{};
|
||
info.struct_size = sizeof(info);
|
||
info.version = YOLO_RUNTIME_INFO_VERSION;
|
||
if (yolo_runtime_get_info(runtime.get(), &info) != YOLO_STATUS_OK) return 6;
|
||
if(dilation_radius && info.task!=YOLO_TASK_SEG) return 2;
|
||
|
||
yolo_image_view_v1_t image{};
|
||
image.struct_size = sizeof(image);
|
||
image.version = YOLO_IMAGE_VIEW_ABI_VERSION;
|
||
image.data = pixels.get();
|
||
image.width = width;
|
||
image.height = height;
|
||
image.row_stride_bytes = width * 3;
|
||
image.pixel_format = YOLO_PIXEL_FORMAT_RGB888;
|
||
yolo_infer_options_v1_t options{};
|
||
options.struct_size = sizeof(options);
|
||
options.version = YOLO_INFER_OPTIONS_VERSION;
|
||
options.seg_geometry_mode = info.task == YOLO_TASK_SEG
|
||
? (geometry_name && std::string_view(geometry_name) == "circle"
|
||
? YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE
|
||
: (geometry_name && std::string_view(geometry_name)=="mean" ? YOLO_SEG_GEOMETRY_MEAN_DIAMETER_CIRCLE
|
||
: (geometry_name && std::string_view(geometry_name)=="fit" ? YOLO_SEG_GEOMETRY_ROBUST_FIT_CIRCLE : YOLO_SEG_GEOMETRY_MIN_AREA_RECT)))
|
||
: YOLO_SEG_GEOMETRY_NONE;
|
||
// ROI 默认数据：W=800（宽）、H=736（高）、L=0（左）、T=0（上），单位为原图像素。
|
||
// 默认初始化同时填写结构体大小、版本并清零保留字段，ABI 5 的字段顺序不变。
|
||
yolo_roi_v1_t roi = yolo_roi_default();
|
||
int roi_values[4] = {0};
|
||
const int roi_mode = parse_example_roi_args(argc, argv, roi_values);
|
||
int has_roi = roi_mode;
|
||
if (roi_mode == 1) {
|
||
// 只提供 L、T：保留上面的默认宽高 800×736，只修改左上角位置。
|
||
roi.x = roi_values[0]; // L：例如 --roi-lt 217 738 中的 217。
|
||
roi.y = roi_values[1]; // T：例如 --roi-lt 217 738 中的 738。
|
||
} else if (roi_mode == 2) {
|
||
// 完整指定 W、H、L、T；例如 --roi-whlt 771 609 217 738。
|
||
roi = yolo_roi_from_whlt(roi_values[0], roi_values[1], roi_values[2], roi_values[3]);
|
||
} else if (roi_mode == 0) {
|
||
// 未传 ROI 参数时沿用模型 roi_info.txt；文件不存在则全图推理。
|
||
// 默认初始化不代表无参数时自动裁剪为 800×736。
|
||
has_roi = load_model_roi(config, width, height, &roi);
|
||
}
|
||
// ROI 必须完整位于原图内；越界、非正宽高由推理接口报错，不自动裁边。
|
||
if (has_roi < 0) {
|
||
fprintf(stderr, "Invalid ROI arguments, roi_info.txt or source image dimensions\n");
|
||
return 7;
|
||
}
|
||
if (has_roi) std::cout << "ROI W=" << roi.width << " H=" << roi.height
|
||
<< " L=" << roi.x << " T=" << roi.y << '\n';
|
||
yolo_result_set_t* raw_results = nullptr;
|
||
error = {};
|
||
error.struct_size = sizeof(error);
|
||
error.version = YOLO_ERROR_INFO_VERSION;
|
||
const auto status =
|
||
has_roi ? yolo_runtime_infer_image_roi(runtime.get(), &image, &roi, &options, &raw_results, &error)
|
||
: yolo_runtime_infer_image(runtime.get(), &image, &options, &raw_results, &error);
|
||
if (status != YOLO_STATUS_OK || !raw_results) {
|
||
std::cerr << "Inference failed: " << error.message << '\n';
|
||
return 7;
|
||
}
|
||
std::unique_ptr<yolo_result_set_t, ResultDeleter> results(raw_results);
|
||
uint32_t count = 0;
|
||
if(yolo_result_set_get_count(results.get(), &count)!=YOLO_STATUS_OK) return 8;
|
||
std::cout << "task=" << info.task << " results=" << count << '\n';
|
||
|
||
// 1. Read performance and typed results directly through the public C ABI.
|
||
yolo_performance_v1_t performance{};
|
||
performance.struct_size=sizeof(performance);performance.version=YOLO_PERFORMANCE_VERSION;
|
||
if(yolo_result_set_get_performance(results.get(),&performance)!=YOLO_STATUS_OK) return 8;
|
||
std::cout << "preprocess_ms=" << performance.preprocess_ms
|
||
<< " inference_ms=" << performance.inference_ms
|
||
<< " postprocess_ms=" << performance.postprocess_ms
|
||
<< " total_ms=" << performance.total_ms << '\n';
|
||
std::vector<yolo_mask_view_v1_t> masks;
|
||
MaskBuffers owned_masks(count);
|
||
const auto region=has_roi?roi:yolo_roi_from_whlt(width,height,0,0);
|
||
for(uint32_t index=0;index<count;++index) {
|
||
yolo_result_common_v1_t common{};
|
||
common.struct_size=sizeof(common);common.version=YOLO_RESULT_COMMON_VERSION;
|
||
if(yolo_result_set_get_common(results.get(),index,&common)!=YOLO_STATUS_OK) return 8;
|
||
// Instance ID is index+1; category/class ID can be shared by many instances.
|
||
std::cout << "instance=" << index+1 << " class=" << common.class_id
|
||
<< " label=" << common.label << " score=" << common.score << '\n';
|
||
if(info.task!=YOLO_TASK_CLS) {
|
||
yolo_bbox_v1_t box{};box.struct_size=sizeof(box);box.version=YOLO_BBOX_VERSION;
|
||
if(yolo_result_set_get_bbox(results.get(),index,&box)!=YOLO_STATUS_OK) return 8;
|
||
std::cout << "bbox=" << box.x << ',' << box.y << ',' << box.width << ',' << box.height << '\n';
|
||
}
|
||
if(info.task==YOLO_TASK_OBB) {
|
||
yolo_rotated_rect_v1_t box{};box.struct_size=sizeof(box);box.version=YOLO_ROTATED_RECT_VERSION;
|
||
if(yolo_result_set_get_obb(results.get(),index,&box)!=YOLO_STATUS_OK) return 8;
|
||
std::cout << "obb center=" << box.center_x << ',' << box.center_y
|
||
<< " size=" << box.width << ',' << box.height << " radians=" << box.angle_radians << '\n';
|
||
}
|
||
if(info.task==YOLO_TASK_POSE) {
|
||
uint32_t keypoint_count=0;
|
||
if(yolo_result_set_get_pose_keypoint_count(results.get(),index,&keypoint_count)!=YOLO_STATUS_OK) return 8;
|
||
std::vector<yolo_keypoint_v1_t> keypoints(keypoint_count);
|
||
if(yolo_result_set_copy_pose_keypoints(results.get(),index,keypoints.data(),keypoint_count,&keypoint_count)!=YOLO_STATUS_OK) return 8;
|
||
for(const auto& point:keypoints)
|
||
std::cout << "keypoint=" << point.x << ',' << point.y << " score=" << point.score << '\n';
|
||
}
|
||
if(info.task==YOLO_TASK_SEG) {
|
||
yolo_seg_geometry_v1_t geometry{};
|
||
geometry.struct_size=sizeof(geometry);geometry.version=YOLO_SEG_GEOMETRY_VERSION;
|
||
if(yolo_result_set_get_seg_geometry(results.get(),index,&geometry)!=YOLO_STATUS_OK) return 8;
|
||
std::cout << "seg type=" << geometry.type << " center=" << geometry.center_x << ',' << geometry.center_y
|
||
<< " radius=" << geometry.radius << " size=" << geometry.width << ',' << geometry.height
|
||
<< " radians=" << geometry.angle_radians << " fit=" << geometry.fit_status << '\n';
|
||
|
||
// 2. Each borrowed binary mask is 0/255, with local dimensions/stride
|
||
// and an origin in the inference input image. Holes are retained.
|
||
yolo_mask_view_v1_t mask{};
|
||
mask.struct_size=sizeof(mask);mask.version=YOLO_MASK_VIEW_VERSION;
|
||
if(yolo_result_set_get_seg_mask(results.get(),index,&mask)!=YOLO_STATUS_OK) return 8;
|
||
// 独立调用公共库；原始推理掩膜和拟合几何保持原值。
|
||
if(dilate_example_mask(&mask,dilation_radius,®ion,&owned_masks.data[index])!=0) return 8;
|
||
masks.push_back(mask);
|
||
std::cout << "binary mask origin=" << mask.origin_x << ',' << mask.origin_y
|
||
<< " size=" << mask.width << ',' << mask.height << " stride=" << mask.row_stride_bytes << '\n';
|
||
}
|
||
}
|
||
|
||
// 3. Every Seg run demonstrates RGB instance IDs in memory by default.
|
||
// --rgb-mask is still accepted for existing scripts, but is no longer needed.
|
||
if(info.task==YOLO_TASK_SEG) {
|
||
const uint64_t stride=uint64_t(width)*3;
|
||
const uint64_t bytes=stride*uint64_t(height);
|
||
if(stride>uint64_t(std::numeric_limits<int32_t>::max()) ||
|
||
bytes>uint64_t(std::numeric_limits<size_t>::max())) return 8;
|
||
std::vector<unsigned char> rgb(static_cast<size_t>(bytes)); // caller owns this buffer
|
||
error={};error.struct_size=sizeof(error);error.version=YOLO_ERROR_INFO_VERSION;
|
||
if(yolo_masks_compose_rgb_ids(masks.data(),count,width,height,rgb.data(),
|
||
static_cast<int32_t>(stride),bytes,&error)!=YOLO_STATUS_OK) {
|
||
std::cerr << "RGB mask failed: " << error.message << '\n';return 8;
|
||
}
|
||
|
||
// 4. Read actual pixel membership: RGB channels encode up to 3 IDs.
|
||
// Background=(0,0,0), single=(id,id,id), double=(id1,id2,id1),
|
||
// triple=(id1,id2,id3). Remove zero/duplicate channels to decode IDs.
|
||
uint64_t foreground=0,overlap_pixels=0;
|
||
std::vector<uint64_t> instance_pixels(count,0);
|
||
for(int y=0;y<height;++y) for(int x=0;x<width;++x) {
|
||
const auto* pixel=rgb.data()+size_t(y)*static_cast<size_t>(stride)+size_t(x)*3;
|
||
unsigned char ids[3]{};int id_count=0;
|
||
for(int channel=0;channel<3;++channel) {
|
||
const auto id=pixel[channel];
|
||
if(!id) continue;
|
||
bool duplicate=false;
|
||
for(int j=0;j<id_count;++j) if(ids[j]==id) duplicate=true;
|
||
if(!duplicate) {
|
||
if(id>count) return 8;
|
||
ids[id_count++]=id;++instance_pixels[id-1];
|
||
}
|
||
}
|
||
if(id_count) {
|
||
if(!foreground) {
|
||
std::cout << "RGB_PIXEL x=" << x << " y=" << y << " RGB="
|
||
<< int(pixel[0]) << ',' << int(pixel[1]) << ',' << int(pixel[2]) << " IDs=";
|
||
for(int j=0;j<id_count;++j) std::cout << (j ? "," : "") << int(ids[j]);
|
||
std::cout << '\n';
|
||
}
|
||
++foreground;if(id_count>1) ++overlap_pixels;
|
||
}
|
||
}
|
||
for(uint32_t i=0;i<count;++i) std::cout << "ID_PIXELS id=" << i+1 << " pixels=" << instance_pixels[i] << '\n';
|
||
std::cout << "RGB_ID_MASK_OK width=" << width << " height=" << height << " instances=" << count
|
||
<< " foreground=" << foreground << " overlaps=" << overlap_pixels << " bytes=" << bytes << '\n';
|
||
// Use rgb here for your processing. No BMP is needed. rgb is freed by
|
||
// vector; borrowed masks stay valid until results is released below.
|
||
}
|
||
std::cout << "MASK_DILATION radius=" << dilation_radius << '\n';
|
||
const int json_status=write_result_json_masks(json_path,config,results.get(),info.task,
|
||
info.task==YOLO_TASK_SEG?masks.data():nullptr,dilation_radius);
|
||
return json_status==0 ? 0 : 8;
|
||
}
|