115 lines
5.6 KiB
C

#include "result_json.h"
#include <stdio.h>
#include <stdlib.h>
static const char* task_name(yolo_task_type_t task) {
switch (task) {
case YOLO_TASK_SEG: return "seg";
case YOLO_TASK_CLS: return "cls";
case YOLO_TASK_OBB: return "obb";
case YOLO_TASK_POSE: return "pose";
default: return "detect";
}
}
static void write_escaped(FILE* file, const char* text) {
for (const unsigned char* cursor = (const unsigned char*)text; cursor && *cursor; ++cursor) {
if (*cursor == '"' || *cursor == '\\') fputc('\\', file);
if (*cursor >= 0x20) fputc(*cursor, file);
}
}
int write_result_json(
const char* path,
const char* model_config,
const yolo_result_set_t* results,
yolo_task_type_t task) {
uint32_t count = 0;
yolo_performance_v1_t performance = {sizeof(performance), YOLO_PERFORMANCE_VERSION};
if (!path || !model_config || !results ||
yolo_result_set_get_count(results, &count) != YOLO_STATUS_OK ||
yolo_result_set_get_performance(results, &performance) != YOLO_STATUS_OK) return -1;
FILE* file = fopen(path, "wb");
if (!file) return -2;
fprintf(file, "{\n \"schemaVersion\": \"2\",\n \"model\": \"");
write_escaped(file, model_config);
fprintf(file, "\",\n \"task\": \"%s\",\n", task_name(task));
fprintf(file, " \"performance\": {\"preprocessMs\": %.3f, \"inferenceMs\": %.3f, \"postprocessMs\": %.3f, \"totalMs\": %.3f},\n",
performance.preprocess_ms, performance.inference_ms,
performance.postprocess_ms, performance.total_ms);
fputs(" \"results\": [", file);
for (uint32_t index = 0; index < count; ++index) {
yolo_result_common_v1_t common = {sizeof(common), YOLO_RESULT_COMMON_VERSION};
yolo_bbox_v1_t bbox = {sizeof(bbox), YOLO_BBOX_VERSION};
if (yolo_result_set_get_common(results, index, &common) != YOLO_STATUS_OK) {
fclose(file); return -3;
}
fprintf(file, "%s\n {\"task\": \"%s\", \"classId\": %d, \"label\": \"",
index ? "," : "", task_name(common.task), common.class_id);
write_escaped(file, common.label);
fprintf(file, "\", \"score\": %.6f", common.score);
if (task != YOLO_TASK_CLS) {
if (yolo_result_set_get_bbox(results, index, &bbox) != YOLO_STATUS_OK) {
fclose(file); return -3;
}
fprintf(file, ", \"bbox\": {\"x\": %.3f, \"y\": %.3f, \"width\": %.3f, \"height\": %.3f",
bbox.x, bbox.y, bbox.width, bbox.height);
if (task == YOLO_TASK_OBB) {
yolo_rotated_rect_v1_t obb = {sizeof(obb), YOLO_ROTATED_RECT_VERSION};
if (yolo_result_set_get_obb(results, index, &obb) != YOLO_STATUS_OK) {
fclose(file); return -3;
}
fprintf(file, ", \"angle\": %.6f", obb.angle_radians);
}
fputc('}', file);
}
if (task == YOLO_TASK_SEG) {
yolo_mask_view_v1_t mask = {sizeof(mask), YOLO_MASK_VIEW_VERSION};
yolo_seg_geometry_v1_t geometry = {sizeof(geometry), YOLO_SEG_GEOMETRY_VERSION};
if (yolo_result_set_get_seg_mask(results, index, &mask) != YOLO_STATUS_OK ||
yolo_result_set_get_seg_geometry(results, index, &geometry) != YOLO_STATUS_OK ||
!mask.data || mask.width <= 0 || mask.height <= 0 || mask.row_stride_bytes < mask.width) {
fclose(file); return -4;
}
fprintf(file, ", \"maskBBox\": {\"x\": %d, \"y\": %d, \"width\": %d, \"height\": %d}",
mask.origin_x, mask.origin_y, mask.width, mask.height);
fprintf(file, ", \"mask\": {\"width\": %d, \"height\": %d, \"strideBytes\": %d, \"originX\": %d, \"originY\": %d}",
mask.width, mask.height, mask.row_stride_bytes, mask.origin_x, mask.origin_y);
fprintf(file, ", \"segGeometry\": {\"type\": \"%s\", \"centerX\": %.3f, \"centerY\": %.3f",
geometry.type == YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE
? "minimumEnclosingCircle" : "minimumAreaRectangle",
geometry.center_x, geometry.center_y);
if (geometry.type == YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE)
fprintf(file, ", \"radius\": %.3f", geometry.radius);
else
fprintf(file, ", \"width\": %.3f, \"height\": %.3f, \"angle\": %.6f",
geometry.width, geometry.height, geometry.angle_radians);
fputc('}', file);
}
if (task == YOLO_TASK_POSE) {
uint32_t keypoint_count = 0, copied = 0;
if (yolo_result_set_get_pose_keypoint_count(results, index, &keypoint_count) != YOLO_STATUS_OK) {
fclose(file); return -5;
}
yolo_keypoint_v1_t* keypoints = keypoint_count
? (yolo_keypoint_v1_t*)malloc(sizeof(*keypoints) * keypoint_count) : NULL;
if (keypoint_count &&
yolo_result_set_copy_pose_keypoints(
results, index, keypoints, keypoint_count, &copied) != YOLO_STATUS_OK) {
free(keypoints); fclose(file); return -5;
}
fputs(", \"keypoints\": [", file);
for (uint32_t point = 0; point < copied; ++point)
fprintf(file, "%s{\"x\": %.3f, \"y\": %.3f, \"score\": %.6f}",
point ? "," : "", keypoints[point].x, keypoints[point].y, keypoints[point].score);
fputc(']', file);
free(keypoints);
}
fputc('}', file);
}
fprintf(file, "%s ]\n}\n", count ? "\n" : "");
fclose(file);
return 0;
}