Untitled

 avatar
unknown
plain_text
a year ago
2.0 kB
4
Indexable
static vsi_status vnn_PostProcessNeuralNetwork
    (
    vsi_nn_graph_t *graph
    )
{
    vsi_status status = vnn_PostProcessOnnx(graph);
    if (status != VSI_SUCCESS)
    {
        return status;
    }

    // Số lượng đầu ra
    const uint32_t num_detections = 3549;
    const uint32_t num_classes = 2;
    const float confidence_threshold = 0.35;

    // Lấy dữ liệu đầu ra từ tensor
    vsi_nn_tensor_t *tensor = vsi_nn_GetTensor(graph, graph->output.tensors[0]);
    uint8_t *tensor_data = NULL;
    vsi_nn_GetTensorHandle(tensor, (void**)&tensor_data);

    if (!tensor_data)
    {
        printf("Error getting tensor data\n");
        return VSI_FAILURE;
    }

    // Dequantize dữ liệu tensor sang float
    uint32_t tensor_size = 1;
    for (uint32_t j = 0; j < tensor->attr.dim_num; j++)
    {
        tensor_size *= tensor->attr.size[j];
    }
    uint32_t stride = vsi_nn_TypeGetBytes(tensor->attr.dtype.vx_type);
    float *output_tensor_buf = calloc(tensor_size, sizeof(float));
    if (!output_tensor_buf)
    {
        printf("Error allocating memory for output tensor buffer\n");
        return VSI_FAILURE;
    }
    for (uint32_t j = 0; j < tensor_size; j++)
    {
        vsi_nn_DtypeToFloat32(&tensor_data[stride * j], &output_tensor_buf[j], &tensor->attr.dtype);
    }

    // Duyệt qua các đầu ra và kiểm tra giá trị confidence
    for (uint32_t i = 0; i < num_detections; i++)
    {
        // Lấy giá trị confidence của hai lớp đầu tiên
        float conf_class0 = output_tensor_buf[i * 34];
        float conf_class1 = output_tensor_buf[i * 34 + 1];

        if (conf_class0 > confidence_threshold)
        {
            printf("fire\n");
        }
        if (conf_class1 > confidence_threshold)
        {
            printf("smoke\n");
        }
    }

    // Giải phóng bộ nhớ
    free(output_tensor_buf);
    vsi_nn_Free(tensor_data);

    return VSI_SUCCESS;
}
Editor is loading...
Leave a Comment