193 lines
10 KiB
C++
193 lines
10 KiB
C++
#include "bmp_loader.h"
|
|
#include "result_json.h"
|
|
#include "yolo_runtime.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); }
|
|
};
|
|
|
|
int main(int argc, char** argv) {
|
|
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]\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;
|
|
|
|
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;
|
|
yolo_result_set_t* raw_results = nullptr;
|
|
error = {};
|
|
error.struct_size = sizeof(error);
|
|
error.version = YOLO_ERROR_INFO_VERSION;
|
|
const auto status =
|
|
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;
|
|
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;
|
|
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.
|
|
}
|
|
const int json_status=write_result_json(json_path,config,results.get(),info.task);
|
|
return json_status==0 ? 0 : 8;
|
|
}
|