master

分支 (1)

管理

管理

master

rv1126-face-recognition
/
rockx_face_detection_rtsp.c

/*
 * 实时检测人脸并通过RTSP传输,VLC播放
 */


#include <assert.h>
#include <fcntl.h>
#include <getopt.h>
#include <pthread.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>

#include "sample_common.h"
#include "rkmedia_api.h"
#include "rkmedia_venc.h"

#include "librtsp/rtsp_demo.h"

#include "rockx.h"

//#define ENABLE_SAVE
#ifdef ENABLE_SAVE
static FILE *g_save_file;    // nv12
static FILE *g_save_file1;   // h264
#endif

static bool quit = false;
static bool readlock = false;

#define ENABLE_RTSP
#ifdef ENABLE_RTSP
/* RTSP相关 */
rtsp_demo_handle g_rtsplive = NULL;
static rtsp_session_handle g_rtsp_session;
#endif

/*!
 * \fn     sigterm_handler
 * \brief  退出信号
 *
 * \param  [in] int sig   #
 *
 * \retval static void
 */
static void sigterm_handler(int sig) {
  fprintf(stderr, "signal %d\n", sig);
  quit = true;
}

/*!
 * \fn     nv12_border
 * \brief  nv12边框
 *
 * \param  [in] char *pic    #
 * \param  [in] int pic_w    #
 * \param  [in] int pic_h    #
 * \param  [in] int rect_x   #
 * \param  [in] int rect_y   #
 * \param  [in] int rect_w   #
 * \param  [in] int rect_h   #
 * \param  [in] int R        #
 * \param  [in] int G        #
 * \param  [in] int B        #
 *
 * \retval int
 */
int nv12_border(char *pic, int pic_w, int pic_h, int rect_x, int rect_y,
                int rect_w, int rect_h, int R, int G, int B) {
  /* Set up the rectangle border size */
  const int border = 5;
  /* RGB convert YUV */
  int Y, U, V;
  Y = 0.299 * R + 0.587 * G + 0.114 * B;
  U = -0.1687 * R + 0.3313 * G + 0.5 * B + 128;
  V = 0.5 * R - 0.4187 * G - 0.0813 * B + 128;
  /* Locking the scope of rectangle border range */
  int j, k;
  for (j = rect_y; j < rect_y + rect_h; j++) {
    for (k = rect_x; k < rect_x + rect_w; k++) {
      if (k < (rect_x + border) || k > (rect_x + rect_w - border) ||
          j < (rect_y + border) || j > (rect_y + rect_h - border)) {
        /* Components of YUV's storage address index */
        int y_index = j * pic_w + k;
        int u_index =
            (y_index / 2 - pic_w / 2 * ((j + 1) / 2)) * 2 + pic_w * pic_h;
        int v_index = u_index + 1;
        /* set up YUV's conponents value of rectangle border */
        pic[y_index] = Y;
        pic[u_index] = U;
        pic[v_index] = V;
      }
    }
  }

  return 0;
}


/*!
 * \fn     GetMediaBuffer
 * \brief  获取nv12原始数据并发送给编码器
 *
 * \param  [in] void *arg   #
 *
 * \retval static void *
 */
static void *GetMediaBuffer(void *arg) {

  pthread_detach(pthread_self());

  int frame_id = 0;
  MEDIA_BUFFER mb = NULL;

  rockx_module_t data_version;
  data_version = ROCKX_MODULE_FACE_DETECTION_V2;

  rockx_ret_t rockx_ret;
  rockx_handle_t face_det_handle;
  rockx_handle_t face_recognize_handle;
  rockx_handle_t face_5landmarks_handle;

  rockx_config_t *config = rockx_create_config();
  rockx_add_config(config, ROCKX_CONFIG_DATA_PATH, "/mnt/nfs/rockx-data-rv1109/");
  rockx_ret = rockx_create(&face_det_handle, data_version, config,
                           sizeof(rockx_config_t));
  if (rockx_ret != ROCKX_RET_SUCCESS) {
    printf("init face_detect error %d\n", rockx_ret);
    return NULL;
  }

  rockx_ret = rockx_create(&face_recognize_handle, ROCKX_MODULE_FACE_RECOGNIZE,
                           config, sizeof(rockx_config_t));
  if (rockx_ret != ROCKX_RET_SUCCESS) {
    printf("init face_recognize error %d\n", rockx_ret);
    return NULL;
  }

  rockx_ret = rockx_create(&face_5landmarks_handle,
                           ROCKX_MODULE_FACE_LANDMARK_5, config, 0);
  if (rockx_ret != ROCKX_RET_SUCCESS) {
    printf("init rockx module ROCKX_MODULE_FACE_LANDMARK_68 error %d\n",
           rockx_ret);
  }

  rockx_image_t input_image;
  input_image.width = 1920;
  input_image.height = 1080;
  input_image.pixel_format = ROCKX_PIXEL_FORMAT_YUV420SP_NV12;

  while (!quit) {

    if(readlock == false)
    {
      mb = RK_MPI_SYS_GetMediaBuffer(RK_ID_VI, 0, -1);
      if (!mb) {
        printf("RK_MPI_SYS_GetMediaBuffer get null buffer!\n");
        break;
      }

      MB_IMAGE_INFO_S stImageInfo = {0};
      int ret = RK_MPI_MB_GetImageInfo(mb, &stImageInfo);
      if (ret)
        printf("Warn: Get image info failed! ret = %d\n", ret);

      //printf("Get Frame:ptr:%p, fd:%d, size:%zu, mode:%d, channel:%d, "
      //     "timestamp:%lld, ImgInfo:<wxh %dx%d, fmt 0x%x>\n",
      //     RK_MPI_MB_GetPtr(mb), RK_MPI_MB_GetFD(mb), RK_MPI_MB_GetSize(mb),
      //     RK_MPI_MB_GetModeID(mb), RK_MPI_MB_GetChannelID(mb),
      //     RK_MPI_MB_GetTimestamp(mb), stImageInfo.u32Width,
      //     stImageInfo.u32Height, stImageInfo.enImgType);

      #ifdef ENABLE_SAVE
        if (g_save_file) {
          fwrite(RK_MPI_MB_GetPtr(mb), 1, RK_MPI_MB_GetSize(mb), g_save_file);
          printf("#Save len-%d to %s\n", RK_MPI_MB_GetSize(mb), g_save_file);
        }
      #endif

      input_image.size = RK_MPI_MB_GetSize(mb);
      input_image.data = (uint8_t *)RK_MPI_MB_GetPtr(mb);

      rockx_object_array_t face_array;
      memset(&face_array, 0, sizeof(face_array));

      rockx_ret =
          rockx_face_detect(face_det_handle, &input_image, &face_array, NULL);
      if (rockx_ret != ROCKX_RET_SUCCESS) {
        printf("rockx_face_detect ERROR %d\n", rockx_ret);
      }

      printf("rockx_face_detect succeed! the face_array count is %d\n", face_array.count);
      if (face_array.count > 0)
      {
        for (int i = 0; i < face_array.count; i++)
        {
          if (1)
          {
            int is_false_face;
            ret = rockx_face_filter(face_5landmarks_handle, &input_image,
                                    &face_array.object[i].box, &is_false_face);
            if (ret != ROCKX_RET_SUCCESS) {
              printf("rockx_face_filter error %d\n", ret);
            }
            // printf("is_false_face: %d\n", is_false_face);
            if (is_false_face)
              continue;
          }

          int left = face_array.object[i].box.left;
          int top = face_array.object[i].box.top;
          int right = face_array.object[i].box.right;
          int bottom = face_array.object[i].box.bottom;
          float score = face_array.object[i].score;

          int w = face_array.object[i].box.right - face_array.object[i].box.left;
          int h = face_array.object[i].box.bottom - face_array.object[i].box.top;

          if (left < 0)
            left = 0;
          if (top < 0)
            top = 0;
          while ((uint32_t)(left + w) >= 1920) {
            w -= 16;
          }
          while ((uint32_t)(top + h) >= 1080) {
            h -= 16;
          }

          nv12_border((char *)input_image.data, input_image.width,
                      input_image.height, left, top, w, h, 0, 0, 255);
        }
      }

      /* 整合数据 */
      memcpy(RK_MPI_MB_GetPtr(mb), input_image.data, input_image.size);

      /* 把数据发给到编码器 */
      RK_MPI_SYS_SendMediaBuffer(RK_ID_VENC, 0, mb);

      RK_MPI_MB_ReleaseBuffer(mb);

      readlock = true;
    }

    usleep(10000);
  }

#ifdef ENABLE_SAVE
  if(g_save_file)
  {
    fclose(g_save_file);
  }
#endif


  return NULL;
}


/*!
 * \fn     OutMediaBuffer
 * \brief  获取编码器数据输出h264
 *
 * \param  [in] void *arg   #
 *
 * \retval static void *
 */
static void *OutMediaBuffer(void *arg)
{
  printf("###############################Start %s thread##########################\n", __func__);

  MEDIA_BUFFER mb = NULL;

  while(!quit)
  {
    if(readlock == true)
    {
      /* 获取编码器里的数据 */
      mb = RK_MPI_SYS_GetMediaBuffer(RK_ID_VENC, 0, -1);
      if (!mb) {
        printf("RK_MPI_SYS_GetMediaBuffer get null buffer!\n");
        break;
      }

      printf("Get packet:ptr:%p, fd:%d, size:%zu, mode:%d, channel:%d, "
             "timestamp:%lld\n",
             RK_MPI_MB_GetPtr(mb), RK_MPI_MB_GetFD(mb), RK_MPI_MB_GetSize(mb),
             RK_MPI_MB_GetModeID(mb), RK_MPI_MB_GetChannelID(mb),
             RK_MPI_MB_GetTimestamp(mb));

      #ifdef ENABLE_SAVE
        if (g_save_file1)
        {
          fwrite(RK_MPI_MB_GetPtr(mb), 1, RK_MPI_MB_GetSize(mb), g_save_file1);
        }
      #endif

      /* 放入对队 */
      const char *nalu_type = "Jpeg data";
      switch(RK_MPI_MB_GetFlag(mb)) {
        case VENC_NALU_IDRSLICE:
          nalu_type = "IDR Slice";
          break;
        case VENC_NALU_PSLICE:
          nalu_type = "P Slice";
          break;
        default:
          nalu_type = "Unknow type";
          break;
      }
      printf("#Write  nalu type: %s, size %zu\n", nalu_type, RK_MPI_MB_GetSize(mb));

	  #ifdef  ENABLE_RTSP
      if (g_rtsplive && g_rtsp_session) {
        rtsp_tx_video(g_rtsp_session, (unsigned char *)RK_MPI_MB_GetPtr(mb), RK_MPI_MB_GetSize(mb), RK_MPI_MB_GetTimestamp(mb));
        rtsp_do_event(g_rtsplive);
      }
      #endif

      RK_MPI_MB_ReleaseBuffer(mb);

      readlock = false;
    }

    usleep(1000);
  }

#ifdef ENABLE_SAVE
  if(g_save_file1)
  {
    fclose(g_save_file1);
  }
#endif

  return NULL;
}


/*!
 * \fn     main
 * \brief  主程序
 *
 * \param  [in] int argc       #
 * \param  [in] char *argv[]   #
 *
 * \retval int
 */
int main(int argc, char *argv[])
{
  RK_U32 u32Width = 1920;
  RK_U32 u32Height = 1080;
  RK_S32 s32CamId = 0;
  RK_CHAR pDeviceName[] = "rkispp_scale0";
  int ret = 0;

  VENC_RC_MODE_E enEncoderMode = VENC_RC_MODE_H264CBR;
  RK_U32 u32Fps = 30;


  signal(SIGINT, sigterm_handler);

#ifdef ENABLE_SAVE
  g_save_file = fopen("output.nv12", "w");
  if (!g_save_file)
    printf("#VENC TEST:: Open ./output.nv12 failed!\n");

  g_save_file1 = fopen("output.h264", "w");
  if (!g_save_file1)
    printf("#VENC TEST:: Open ./output.h264 failed!\n");

#endif

  #ifdef ENABLE_RTSP
    // init rtsp
    g_rtsplive = create_rtsp_demo(554);
    g_rtsp_session = rtsp_new_session(g_rtsplive, "/live/main_stream");

    rtsp_set_video(g_rtsp_session, RTSP_CODEC_ID_VIDEO_H264, NULL, 0);

    rtsp_sync_video_ts(g_rtsp_session, rtsp_get_reltime(), rtsp_get_ntptime());
  #endif

  RK_MPI_SYS_Init();

  VI_CHN_ATTR_S vi_chn_attr;
  vi_chn_attr.pcVideoNode = pDeviceName;
  vi_chn_attr.u32BufCnt = 3;
  vi_chn_attr.u32Width = u32Width;
  vi_chn_attr.u32Height = u32Height;
  vi_chn_attr.enPixFmt = IMAGE_TYPE_NV12;
  vi_chn_attr.enWorkMode = VI_WORK_MODE_NORMAL;
  vi_chn_attr.enBufType = VI_CHN_BUF_TYPE_MMAP;
  ret = RK_MPI_VI_SetChnAttr(s32CamId, 0, &vi_chn_attr);
  ret |= RK_MPI_VI_EnableChn(s32CamId, 0);
  if (ret) {
    printf("Create VI[0] failed! ret=%d\n", ret);
    return -1;
  }

  /* 编码初始化 */
  VENC_CHN_ATTR_S venc_chn_attr;
  memset(&venc_chn_attr, 0, sizeof(VENC_CHN_ATTR_S));
  venc_chn_attr.stVencAttr.u32PicWidth = u32Width;
  venc_chn_attr.stVencAttr.u32PicHeight = u32Height;
  venc_chn_attr.stVencAttr.u32VirWidth = u32Width;
  venc_chn_attr.stVencAttr.u32VirHeight = u32Height;
  venc_chn_attr.stVencAttr.imageType = IMAGE_TYPE_NV12;

  venc_chn_attr.stVencAttr.enType = RK_CODEC_TYPE_H264;
  venc_chn_attr.stVencAttr.u32Profile = 77;
  venc_chn_attr.stRcAttr.enRcMode = VENC_RC_MODE_H264CBR;
  venc_chn_attr.stRcAttr.stH264Cbr.u32Gop = 2 * u32Fps;
  venc_chn_attr.stRcAttr.stH264Cbr.u32BitRate = u32Width * u32Height;
  venc_chn_attr.stRcAttr.stH264Cbr.fr32DstFrameRateDen = 1;
  venc_chn_attr.stRcAttr.stH264Cbr.fr32DstFrameRateNum = u32Fps;
  venc_chn_attr.stRcAttr.stH264Cbr.u32SrcFrameRateDen = 1;
  venc_chn_attr.stRcAttr.stH264Cbr.u32SrcFrameRateNum = u32Fps;

  ret = RK_MPI_VENC_CreateChn(0, &venc_chn_attr);
  if (ret) {
    printf("ERROR: Create venc failed!\n");
    exit(0);
  }

  printf("====> %s initial finish\n", __func__);


  pthread_t out_thread;
  pthread_create(&out_thread, NULL, OutMediaBuffer, NULL);

  pthread_t read_thread;
  pthread_create(&read_thread, NULL, GetMediaBuffer, NULL);


  ret = RK_MPI_VI_StartStream(s32CamId, 0);
  if (ret) {
    printf("Start VI[0] failed! ret=%d\n", ret);
    return -1;
  }

  while (!quit) {
    usleep(500000);
  }

#ifdef ENABLE_RTSP
  if (g_rtsplive)
    rtsp_del_demo(g_rtsplive);
#endif


  RK_MPI_VI_DisableChn(s32CamId, 0);

  RK_MPI_VENC_DestroyChn(0);

  printf("%s exit!\n", __func__);
  return 0;
}