67 lines
2.8 KiB
C

#include "bmp_loader.h"
#include "result_json.h"
#include "yolo_runtime.h"
#include <stdio.h>
#include <string.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) {
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]\n", argv[0]);
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;
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 : YOLO_SEG_GEOMETRY_MIN_AREA_RECT)
: YOLO_SEG_GEOMETRY_NONE;
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(runtime, &image, &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);
const int json_status = write_result_json(json_path, config, results, info.task);
yolo_result_set_release(results);
yolo_runtime_destroy(runtime);
free_bmp_rgb(pixels);
return json_status == 0 ? 0 : 8;
}