yolov5-inferrence/main.cpp

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>
#include "yolov5.h"
#include <sys/stat.h>
#include <sys/types.h>
#include "vertix.h"

/*Include ffmpeg header file*/
#ifdef __cplusplus
extern "C"
{
#endif

#include <libavformat/avformat.h>

#ifdef __cplusplus
};
#endif

#include "MppDecode.h"
#include "rga.h"
#include "pose_body.h"
#include "rknn.h"

long max_time_elapsed = 0;
char *server = NULL;

void deInit(MppPacket *packet, MppFrame *frame, MppCtx ctx, char *buf, MpiDecLoopData data)
{
    if (packet)
    {
        mpp_packet_deinit(packet);
        packet = NULL;
    }

    if (frame)
    {
        mpp_frame_deinit(frame);
        frame = NULL;
    }

    if (ctx)
    {
        mpp_destroy(ctx);
        ctx = NULL;
    }

    if (buf)
    {
        mpp_free(buf);
        buf = NULL;
    }

    if (data.pkt_grp)
    {
        mpp_buffer_group_put(data.pkt_grp);
        data.pkt_grp = NULL;
    }

    if (data.frm_grp)
    {
        mpp_buffer_group_put(data.frm_grp);
        data.frm_grp = NULL;
    }

    if (data.fp_output)
    {
        fclose(data.fp_output);
        data.fp_output = NULL;
    }

    if (data.fp_input)
    {
        fclose(data.fp_input);
        data.fp_input = NULL;
    }
}

void usage(char *name)
{
    printf("\nUsage: %s [-s $rtsp-stream-address] \n\t -s specify the rtsp address, like \"rtsp://admin:$password@$ip:554/h264/ch1/main/av_stream\"\n", name);
}

void example_server(char *name)
{
    printf("\nexample server: rtsp://admin:JEGSSI@172.28.1.208:554/h264/ch1/main/av_stream\n");
}

static const char *help_info[100] = {
    "rtsp server URL",
    "show example server",
    // "model name from /usr/local/models/",
    "model name from ./models/",
    "camera id",
    "type id",
    "upload port",
    "workspace, use comma spearated coordinate, like 0,1,2,3"
    "threshold, like 0.7, default is 0.6"};

static struct option long_options[] = {
    {"server", required_argument, NULL, 's'},
    {"example", no_argument, NULL, 'e'},
    {"model", required_argument, NULL, 'm'},
    {"camera-id", required_argument, NULL, 'c'},
    {"type-id", required_argument, NULL, 't'},
    {"upload-port", required_argument, NULL, 'p'},
    {"workspace", required_argument, NULL, 'w'},
    {"threshold", required_argument, NULL, 'r'},
    {"time-period", required_argument, NULL, 'P'},
    {0, 0, 0}};

void print_usage(char *cmd)
{
    printf("Usage: %s [OPTION]\n\nGetting help:\n", cmd);
    struct option *opt;
    for (int i = 0;; i++)
    {
        opt = long_options + i;
        if (opt->name == NULL)
        {
            break;
        }
        printf("  -%c, -%s\t%s\n", opt->val, opt->name, help_info[i]);
    }
}

float *parse_workspace_point(char *input, int maxnumber, int *size)
{
    printf("got workspace points: %s\n", input);
    float *result = (float *)malloc(sizeof(float) * maxnumber);
    int count = 0;
    char *token = strtok(input, ",");
    while (token != NULL && count < maxnumber)
    {
        printf("token > %s\n", token);
        int k = atoi(token);
        printf("converted to int is %d\n", k);
        result[count++] = float(k);
        printf("stored is %f\n", result[count - 1]);
        token = strtok(NULL, ",");
    }
    *size = count;
    return result;
}

void parse_options(int argc, char **argv)
{
    int c;

    bool polygon_inited = false;
    char *workspace_points = (char*)"";
    float *points;
    int thresh;
    conf_threshold = DEFAULT_CONF_THRESHOLD;
    char *cmd = argv[0];
    int ok;
    while (1)
    {
        int option_index = 0;
        c = getopt_long_only(argc, argv, "es:m:w:c:t:p:r:P:", long_options, &option_index);
        if (c == -1)
        {
            break;
        }
        switch (c)
        {
        case 'c':
            CameraID = optarg;
            break;
        case 't':
            TypeID = atoi(optarg);
            break;
        case 'p':
            UploadPort = atoi(optarg);
            break;
        case 's':
            server = optarg;
            break;
        case 'e':
            example_server(cmd);
            exit(0);
        case 'w':
            workspace_points = optarg;
            int workspace_size;
            points = parse_workspace_point(workspace_points, 50, &workspace_size);

            init_polygon(points, workspace_size);
            polygon_inited = true;
            free(points);
            break;
        case 'm':
            global_model = optarg;
            break;
        case 'r':
            thresh = atoi(optarg);
            if (thresh < 30) {
                thresh = 30;
            } else if (thresh > 90) {
                thresh = 90;
            }
            conf_threshold = (float)thresh / 100.0;
            break;
        case 'P':
            ok = get_all_time_period_from_arguments(optarg);
            if (ok != 0) {
                printf("failed to parse time period\n");
                exit(-1);
            }
            break;
        default:
            print_usage(cmd);
            exit(-1);
        }
    }

    if (!polygon_inited) {
        init_polygon(NULL, 0);
    }
}

int main(int argc, char *argv[])
{
    int o;
    // const char *optstring = "es:m:h:u:p:c:";

    /*
        char tempbuf[20];
        for (int i = 0;; i++) {
            cJSON *obj = cJSON_CreateObject();
            cJSON_AddNumberToObject(obj, "index", i);
            cJSON_AddStringToObject(obj, "name", "hello");
            publishCJSONInfo("sensor/a", obj, 0);
            cJSON_Delete(obj);
            sleep(1);
        }
    */

    // char *server = "rtsp://admin:JEGSSI@172.28.1.208:554/h264/ch1/main/av_stream";
    char *tempm = (char *)malloc(sizeof(char) * 128);

    parse_options(argc, argv);
    if (CameraID == NULL)
    {
        printf("no camera id specified");
        return -1;
    }
    if (global_model == NULL)
    {
        global_model = (char*)"yolov5s";
    }

    char targetdir[120];
    mkdir("/usr");
    mkdir("/usr/data");
    mkdir("/usr/data/camera");
    printf("mkdir camera");
    snprintf(targetdir, 120, "/usr/data/camera/%s", CameraID);
    mkdir(targetdir);
    snprintf(targetdir, 120, "/usr/data/camera/%s/picsave", CameraID);
    mkdir(targetdir);
    snprintf(targetdir, 120, "/usr/data/camera/%s/frames", CameraID);
    char rmdir[200];
    snprintf(rmdir, 200, "rm -rf %s", targetdir);
    system(rmdir);
    mkdir(targetdir);

    // initMQTTClient();
    /*
    if (global_model == NULL)
    {
        sprintf(tempm, "%s", "coco");
        global_model = tempm;
    }
    */

    // do_init_model();

    int ret;
    memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t));

    init_post_process();

    char model_path[128];
    snprintf(model_path, 127, "%s/%s.rknn", MODELDIR, global_model);

    ret = init_yolov5_model(model_path, &rknn_app_ctx, print_object_detect_result);
    if (ret != 0)
    {
        printf("init_yolov5_model fail! ret=%d model_path=%s\n", ret, model_path);

        deinit_post_process();

        ret = release_yolov5_model(&rknn_app_ctx);
        if (ret != 0)
        {
            printf("release_yolov5_model fail! ret=%d\n", ret);
        }
        return -1;
    }

    printf("server: %s\n", server);

    // initRkRga();
    // if (init_pose_body_context() != 0)
    // {
    //     printf("failed to init pose_body context");
    //     return -1;
    // }
    AVFormatContext *pFormatCtx = NULL;
    AVDictionary *options = NULL;
    AVPacket *av_packet = NULL;
    // char filepath[] = "rtsp://admin:openiot123456@192.168.1.64:554/h264/ch1/main/av_stream";// rtsp 地址
    // char filepath[] = "rtsp://admin:JEGSSI@172.28.1.208:554/h264/ch1/main/av_stream";
    // char filepath[] = "rtsp://admin:openiot123456@192.168.1.64:554/h264/ch39/sub/av_stream";// rtsp 地址
    // char filepath[] = "rtsp://192.168.180.1:8554/cam";// rtsp 地址
    // char filepath[] = "/root/test.h264";// rtsp 地址

    // av_register_all(); // 函数在ffmpeg4.0以上版本已经被废弃，所以4.0以下版本就需要注册初始函数
    avformat_network_init();
    // av_dict_set(&options, "buffer_size", "1024000", 0); //设置缓存大小,1080p可将值跳到最大
    av_dict_set(&options, "buffer_size", "10240000", 0); // 设置缓存大小,1080p可将值跳到最大
    av_dict_set(&options, "rtsp_transport", "tcp", 0);   // 以tcp的方式打开,
    av_dict_set(&options, "stimeout", "5000000", 0);     // 设置超时断开链接时间，单位us
    av_dict_set(&options, "max_delay", "500000", 0);     // 设置最大时延

    // pFormatCtx = avformat_alloc_context(); // 用来申请AVFormatContext类型变量并初始化默认参数,申请的空间
    pFormatCtx = avformat_alloc_context(); // 用来申请AVFormatContext类型变量并初始化默认参数,申请的空间

    // 打开网络流或文件流
    if (avformat_open_input(&pFormatCtx, server, NULL, &options) != 0)
    {
        printf("Couldn't open input stream.\n");
        return 0;
    }

    // 获取视频文件信息
    if (avformat_find_stream_info(pFormatCtx, NULL) < 0)
    {
        printf("Couldn't find stream information.\n");
        return 0;
    }

    // 查找码流中是否有视频流
    int videoindex = -1;
    unsigned i = 0;
    for (i = 0; i < pFormatCtx->nb_streams; i++)
        if (pFormatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            videoindex = i;
            break;
        }
    if (videoindex == -1)
    {
        printf("Didn't find a video stream.\n");
        return 0;
    }

    av_packet = (AVPacket *)av_malloc(sizeof(AVPacket)); // 申请空间，存放的每一帧数据 （h264、h265）

    //// 初始化
    MPP_RET ret2 = MPP_OK;
    size_t file_size = 0;

    // base flow context
    MppCtx ctx = NULL;
    MppApi *mpi = NULL;

    // input / output
    MppPacket packet = NULL;
    MppFrame frame = NULL;

    MpiCmd mpi_cmd = MPP_CMD_BASE;
    MppParam param = NULL;
    RK_U32 need_split = 1;
    //    MppPollType timeout = 5;

    // paramter for resource malloc
    // RK_U32 width        = 2560;
    // RK_U32 height       = 1440;
    RK_U32 width = 1080;
    RK_U32 height = 768;
    MppCodingType type = MPP_VIDEO_CodingAVC;

    // resources
    char *buf = NULL;
    size_t packet_size = 8 * 1024 * 1024;
    MppBuffer pkt_buf = NULL;
    MppBuffer frm_buf = NULL;

    MpiDecLoopData data;

    mpp_log("mpi_dec_test start\n");
    memset(&data, 0, sizeof(data));

    /*
    data.fp_output = fopen("./tenoutput.yuv", "w+");
    if (NULL == data.fp_output) {
        mpp_err("failed to open output file %s\n", "tenoutput.yuv");
        deInit(&packet, &frame, ctx, buf, data);
    }
    */

    mpp_log("mpi_dec_test decoder test start w %d h %d type %d\n", width, height, type);

    // decoder demo
    ret2 = mpp_create(&ctx, &mpi);

    if (MPP_OK != ret2)
    {
        mpp_err("mpp_create failed\n");
        deInit(&packet, &frame, ctx, buf, data);
    }

    // NOTE: decoder split mode need to be set before init
    mpi_cmd = MPP_DEC_SET_PARSER_SPLIT_MODE;
    param = &need_split;
    ret2 = mpi->control(ctx, mpi_cmd, param);
    if (MPP_OK != ret2)
    {
        mpp_err("mpi->control failed\n");
        deInit(&packet, &frame, ctx, buf, data);
    }

    mpi_cmd = MPP_SET_INPUT_BLOCK;
    param = &need_split;
    ret2 = mpi->control(ctx, mpi_cmd, param);
    if (MPP_OK != ret2)
    {
        mpp_err("mpi->control failed\n");
        deInit(&packet, &frame, ctx, buf, data);
    }

    ret2 = mpp_init(ctx, MPP_CTX_DEC, type);
    if (MPP_OK != ret2)
    {
        mpp_err("mpp_init failed\n");
        deInit(&packet, &frame, ctx, buf, data);
    }

    data.ctx = ctx;
    data.mpi = mpi;
    data.eos = 0;
    data.packet_size = packet_size;
    data.frame = frame;
    data.frame_count = 0;

    int count = 0;

    current_time_gap_t gap;
    char outfilename[128];

    get_current_time_gap(&gap);
    time_t next = gap.next_gap;

    snprintf(outfilename, 128, "/usr/data/camera/%s/%ld.h264", CameraID, convert_time_to_name(gap.now_gap));

    FILE *fpSave = fopen(outfilename, "ab");

    int errorcount = 0;

    // 这边可以调整i的大小来改变文件中的视频时间,每 +1 就是一帧数据

    struct timeval tpend1, tpend2;
    long usec1 = 0;
    gettimeofday(&tpend1, NULL);

    for (;;)
    {
        if (errorcount > 50)
        {
            printf("error count expires 50(%d), exitting...\n", errorcount);
            break;
        }
        // printf("try reading...\n");
        if (av_read_frame(pFormatCtx, av_packet) >= 0)
        {
            printf("errorcount set to 0");
            errorcount = 0;
            if (av_packet->stream_index == videoindex)
            {
                // mpp_log("--------------\ndata size is: %d\n-------------", av_packet->size);

                count++;
                printf("handling [%d]th frame\n", count);
                decode_simple(&data, av_packet);

                if (fpSave != NULL)
                {
                    fwrite(av_packet->data, av_packet->size, 1, fpSave);
                    time_t now = time(NULL);
                    if (now >= next)
                    {
                        fclose(fpSave);
                        get_current_time_gap(&gap);
                        next = gap.next_gap;

                        snprintf(outfilename, 128, "/usr/data/camera/%s/%ld.h264", CameraID, convert_time_to_name(gap.now_gap));
                        // snprintf(outfilename, 20, "%ld.h264", convert_time_to_name(gap.now_gap));
                        fpSave = fopen(outfilename, "ab");
                    }
                }
            }
            if (av_packet != NULL)
                av_packet_unref(av_packet);
            // mpp_log("%d", i);
        }
        else
        {
            errorcount++;
        }
    }

    printf("max time elapsed: %f\n", (float)max_time_elapsed / CLOCKS_PER_SEC * 1000);
    if (fpSave != NULL)
    {
        fclose(fpSave);
        fpSave = NULL;
    }
    deInit(&packet, &frame, ctx, buf, data);
    av_free(av_packet);
    avformat_close_input(&pFormatCtx);
    printf("done\n");

    deinit_post_process();

    release_yolov5_model(&rknn_app_ctx);
}