161 lines
8.5 KiB
C
161 lines
8.5 KiB
C
#include "bmp_loader.h"
|
||
#include "result_json.h"
|
||
#include "yolo_runtime.h"
|
||
#include "mask_dilation.h"
|
||
|
||
#include <stdio.h>
|
||
#include <string.h>
|
||
#include <stdlib.h>
|
||
#include <limits.h>
|
||
|
||
static const char* value_after(int argc, char** argv, const char* option) {
|
||
for (int index = 1; index + 1 < argc; ++index)
|
||
if (strcmp(argv[index], option) == 0) return argv[index + 1];
|
||
return NULL;
|
||
}
|
||
|
||
int main(int argc, char** argv) {
|
||
int dilation_radius=0;
|
||
if(parse_mask_dilation(argc,argv,&dilation_radius)!=0) {fprintf(stderr,"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) {
|
||
fprintf(stderr, "Usage: %s --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", argv[0]);
|
||
return 2;
|
||
}
|
||
if(geometry_name && strcmp(geometry_name,"rect") && strcmp(geometry_name,"circle") &&
|
||
strcmp(geometry_name,"mean") && strcmp(geometry_name,"fit")) return 2;
|
||
if (yolo_runtime_get_abi_version() != YOLO_RUNTIME_ABI_VERSION) return 3;
|
||
unsigned char* pixels = NULL;
|
||
int width = 0, height = 0;
|
||
if (load_bmp_rgb(image_path, &pixels, &width, &height) != 0) return 4;
|
||
|
||
yolo_error_info_v1_t error = {sizeof(error), YOLO_ERROR_INFO_VERSION};
|
||
yolo_runtime_t* runtime = NULL;
|
||
if (yolo_runtime_create(config, &runtime, &error) != YOLO_STATUS_OK || !runtime) {
|
||
fprintf(stderr, "Create failed: %s\n", error.message);
|
||
free_bmp_rgb(pixels);
|
||
return 5;
|
||
}
|
||
yolo_runtime_info_v1_t info = {sizeof(info), YOLO_RUNTIME_INFO_VERSION};
|
||
if (yolo_runtime_get_info(runtime, &info) != YOLO_STATUS_OK) return 6;
|
||
if(dilation_radius && info.task!=YOLO_TASK_SEG) {yolo_runtime_destroy(runtime);free_bmp_rgb(pixels);return 2;}
|
||
yolo_image_view_v1_t image = {sizeof(image), YOLO_IMAGE_VIEW_ABI_VERSION};
|
||
image.data = pixels;
|
||
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 = {sizeof(options), YOLO_INFER_OPTIONS_VERSION};
|
||
options.seg_geometry_mode = info.task == YOLO_TASK_SEG
|
||
? (geometry_name && strcmp(geometry_name, "circle") == 0
|
||
? YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE
|
||
: (geometry_name && strcmp(geometry_name,"mean")==0 ? YOLO_SEG_GEOMETRY_MEAN_DIAMETER_CIRCLE
|
||
: (geometry_name && strcmp(geometry_name,"fit")==0 ? 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");
|
||
yolo_runtime_destroy(runtime); free_bmp_rgb(pixels);
|
||
return 7;
|
||
}
|
||
if (has_roi) printf("ROI W=%d H=%d L=%d T=%d\n", roi.width, roi.height, roi.x, roi.y);
|
||
yolo_result_set_t* results = NULL;
|
||
error = (yolo_error_info_v1_t){sizeof(error), YOLO_ERROR_INFO_VERSION};
|
||
const yolo_status_t status =
|
||
yolo_runtime_infer_image_roi(runtime, &image, has_roi ? &roi : NULL, &options, &results, &error);
|
||
if (status != YOLO_STATUS_OK || !results) {
|
||
fprintf(stderr, "Inference failed: %s\n", error.message);
|
||
yolo_runtime_destroy(runtime);
|
||
free_bmp_rgb(pixels);
|
||
return 7;
|
||
}
|
||
uint32_t count = 0;
|
||
yolo_result_set_get_count(results, &count);
|
||
printf("task=%d results=%u\n", info.task, count);
|
||
yolo_mask_view_v1_t masks[255];
|
||
unsigned char* owned_masks[255]={0};
|
||
memset(masks,0,sizeof(masks));
|
||
const yolo_roi_v1_t region=has_roi?roi:yolo_roi_from_whlt(width,height,0,0);
|
||
int mask_status=0;
|
||
/* Seg always demonstrates in-memory RGB IDs; --rgb-mask remains accepted. */
|
||
if(info.task==YOLO_TASK_SEG) {
|
||
|
||
if(count>255) {fprintf(stderr,"RGB supports at most 255 instances\n");mask_status=-1;}
|
||
for(uint32_t i=0;i<count && mask_status==0;++i) {
|
||
masks[i].struct_size=sizeof(masks[i]);masks[i].version=YOLO_MASK_VIEW_VERSION;
|
||
/* Local 0/255 data is borrowed until results is released. */
|
||
if(yolo_result_set_get_seg_mask(results,i,&masks[i])!=YOLO_STATUS_OK) mask_status=-1;
|
||
/* 独立调用公共库膨胀；输出由调用方持有，保留实例重叠。 */
|
||
if(mask_status==0 && dilate_example_mask(&masks[i],dilation_radius,®ion,&owned_masks[i])!=0) mask_status=-1;
|
||
}
|
||
const uint64_t stride=(uint64_t)width*3,bytes=stride*(uint64_t)height;
|
||
unsigned char* rgb=NULL;
|
||
if(stride>INT_MAX || bytes>SIZE_MAX) mask_status=-1;
|
||
if(mask_status==0) {
|
||
rgb=(unsigned char*)malloc((size_t)bytes);
|
||
if(!rgb) mask_status=-1;
|
||
}
|
||
if(mask_status==0) {
|
||
error=(yolo_error_info_v1_t){sizeof(error),YOLO_ERROR_INFO_VERSION};
|
||
if(yolo_masks_compose_rgb_ids(masks,count,width,height,rgb,(int32_t)stride,bytes,&error)!=YOLO_STATUS_OK) {
|
||
fprintf(stderr,"RGB mask failed: %s\n",error.message);mask_status=-1;
|
||
} else {
|
||
uint64_t foreground=0,instance_pixels[256]={0};
|
||
for(int y=0;y<height;++y) for(int x=0;x<width;++x) {
|
||
const unsigned char* pixel=rgb+(size_t)y*(size_t)stride+(size_t)x*3;
|
||
unsigned char ids[3]={0};int n=0;
|
||
/* Remove zero and duplicate channel values to decode membership. */
|
||
for(int c=0;c<3;++c) {
|
||
unsigned char id=pixel[c];int duplicate=0;
|
||
if(!id) continue;
|
||
for(int j=0;j<n;++j) if(ids[j]==id) duplicate=1;
|
||
if(!duplicate) {ids[n++]=id;++instance_pixels[id];}
|
||
}
|
||
if(n) {
|
||
if(!foreground) {
|
||
printf("RGB_PIXEL x=%d y=%d RGB=%u,%u,%u IDs=",x,y,pixel[0],pixel[1],pixel[2]);
|
||
for(int j=0;j<n;++j) printf("%s%u",j?",":"",ids[j]);
|
||
printf("\n");
|
||
}
|
||
++foreground;
|
||
}
|
||
}
|
||
for(uint32_t i=1;i<=count;++i) printf("ID_PIXELS id=%u pixels=%llu\n",i,(unsigned long long)instance_pixels[i]);
|
||
printf("RGB_ID_MASK_OK width=%d height=%d instances=%u foreground=%llu bytes=%llu\n",
|
||
width,height,count,(unsigned long long)foreground,(unsigned long long)bytes);
|
||
/* Process rgb here directly; saving a bitmap is optional. */
|
||
}
|
||
}
|
||
free(rgb); /* Caller owns RGB; result set separately owns binary masks. */
|
||
}
|
||
const int json_status=mask_status==0 ? write_result_json_masks(json_path,config,results,info.task,
|
||
info.task==YOLO_TASK_SEG?masks:NULL,dilation_radius) : -1;
|
||
for(int i=0;i<255;++i) free(owned_masks[i]);
|
||
printf("MASK_DILATION radius=%d\n",dilation_radius);
|
||
yolo_result_set_release(results);
|
||
yolo_runtime_destroy(runtime);
|
||
free_bmp_rgb(pixels);
|
||
return json_status == 0 && mask_status == 0 ? 0 : 8;
|
||
}
|