int main()

in VirtualStage/KinectMaskGenerator/main.cpp [72:333]


int main(int argc, char *argv[])
{
  if (argc < 5)
  {
    printf("USAGE: VIDEO_INPUT_PATH MASK_OUTPUT_PATH START DURATION [COLOR_IMG_SUFFIX]");
    exit(1);
  }

  char *INPUT = argv[1];
  char *OUTPUT = argv[2];

  uint64_t timestamp_start_us = (uint64_t)std::stoi(argv[3]) * (uint64_t)1e6;
  printf("Starting at second %f\n", timestamp_start_us / 1e6);
  int duration = std::stoi(argv[4]);
  uint64_t timestamp_end_us = 0;
  if (duration > 0)
  {
    timestamp_end_us = (uint64_t)std::stoi(argv[4]) * (uint64_t)1e6 + timestamp_start_us;
    printf("Finishing at second %f\n", timestamp_end_us / 1e6);
  }

  bool write_color = false;
  char *color_img_suffix = "";
  if (argc >= 6)
  {
    write_color = true;
    color_img_suffix = argv[5];
  }

  SyncronizationFileWriter syncFileWriter;
  syncFileWriter.init(OUTPUT);

  k4a_playback_t playback_handle = NULL;
  VERIFY(k4a_playback_open(INPUT, &playback_handle), "Failed to open recording\n");

  uint64_t recording_length = k4a_playback_get_recording_length_usec(playback_handle);
  printf("Recording is %lld seconds long\n", recording_length / 1000000);

  k4a_calibration_t sensor_calibration;
  VERIFY(k4a_playback_get_calibration(playback_handle, &sensor_calibration),
         "Get camera calibration failed!");

  k4a_playback_set_color_conversion(playback_handle,
                                    K4A_IMAGE_FORMAT_COLOR_BGRA32);

  // Create the tracker
  k4abt_tracker_t tracker = NULL;
  k4abt_tracker_configuration_t tracker_config = K4ABT_TRACKER_CONFIG_DEFAULT;
  VERIFY(k4abt_tracker_create(&sensor_calibration, tracker_config, &tracker), "Body tracker initialization failed!");

  int frame_count = 0;
  char buffer[120];
  int color_image_width_pixels = sensor_calibration.color_camera_calibration.resolution_width;
  int color_image_height_pixels = sensor_calibration.color_camera_calibration.resolution_height;
  int depth_image_width_pixels = sensor_calibration.depth_camera_calibration.resolution_width;
  int depth_image_height_pixels = sensor_calibration.depth_camera_calibration.resolution_height;
  k4a_image_t mask_image = NULL;
  k4a_image_t transformed_mask_image = NULL;
  k4a_image_t transformed_depth_image = NULL;
  // transform depth camera into color camera geometry
  VERIFY(k4a_image_create(K4A_IMAGE_FORMAT_CUSTOM8,
                          color_image_width_pixels,
                          color_image_height_pixels,
                          color_image_width_pixels * (int)sizeof(uint8_t),
                          &transformed_mask_image),
         "Could not create image for mask transformation!");

  VERIFY(k4a_image_create(K4A_IMAGE_FORMAT_DEPTH16,
                          color_image_width_pixels,
                          color_image_height_pixels,
                          color_image_width_pixels * (int)sizeof(uint16_t),
                          &transformed_depth_image),
         "Could not create image for depth transformation!");
  k4a_transformation_t transformation = k4a_transformation_create(&sensor_calibration);
  Mat last_mask = cv::Mat::zeros(cv::Size(1920, 1080), CV_8U);
  bool writing = false;
  uint64_t initial_duration = 0;

  while (true)
  {
    frame_count++;

    // Get a capture
    k4a_capture_t sensor_capture;
    k4a_stream_result_t get_capture_result = k4a_playback_get_next_capture(playback_handle, &sensor_capture);
    if (get_capture_result == K4A_STREAM_RESULT_EOF)
    {
      break;
    }
    else if (get_capture_result != K4A_STREAM_RESULT_SUCCEEDED)
    {
      if (writing)
      {
        // If the frame is corrupt, we write last interval timestamp and last mask
        printf("Warning: Get capture returned error %d on frame %d\n", get_capture_result, frame_count);
        snprintf(buffer, 128, "%s/%04d_masksAK.png", OUTPUT, frame_count);
        cv::imwrite(buffer, last_mask);
        syncFileWriter.writeStamp(frame_count);
      }
      continue;
    }

    // if there is no color image, we just consider that frame didn't exist
    k4a_image_t color_image = k4a_capture_get_color_image(sensor_capture);
    if (color_image == 0)
    {
      printf("Warning: No color image %d on frame %d, dropped\n", get_capture_result, frame_count);
      frame_count--;
      k4a_capture_release(sensor_capture);
      continue;
    }

    // write color image
    if (writing && write_color)
    {
      Mat color = color_to_opencv(color_image);
      snprintf(buffer, 128, "%s/%04d_%s.png", OUTPUT, frame_count, color_img_suffix);
      cv::imwrite(buffer, color);
    }

    // get frame timestamp information
    uint64_t duration_us = k4a_image_get_device_timestamp_usec(color_image);
    k4a_image_release(color_image);
    if (initial_duration == 0)
    { // get the first duration to compare with the rest
      initial_duration = duration_us;
    }
    current_interval = duration_us - last_duration;
    last_duration = duration_us;
    // sometimes first timestamp is wrong, in that case we use default
    if (current_interval / 1e6 > 5)
    {
      current_interval = default_interval;
    }

    // check if we are in the interested piece of video
    if (!writing && (duration_us - initial_duration) >= timestamp_start_us)
    {
      writing = true; // start writing
      frame_count = 1;
    }
    else if (writing && duration > 0 && (duration_us - initial_duration) >= timestamp_end_us)
    {
      break; // done
    }

    // don't queue if we are to too far from the start
    if (!writing && (duration_us - initial_duration) < (timestamp_start_us - int(1e6)))
    {
      k4a_capture_release(sensor_capture);
      continue;
    }

    // queque capture
    k4a_wait_result_t queue_capture_result = k4abt_tracker_enqueue_capture(tracker, sensor_capture, K4A_WAIT_INFINITE);
    k4a_capture_release(sensor_capture);
    if (queue_capture_result == K4A_WAIT_RESULT_TIMEOUT)
    {
      // It should never hit timeout when K4A_WAIT_INFINITE is set.
      printf("Error! Add capture to tracker process queue timeout!\n");
      exit(1);
    }
    else if (queue_capture_result == K4A_WAIT_RESULT_FAILED)
    {
      if (writing)
      {
        // If for some reason (probably there is no depth information) we don't have a mask for
        //   this frame, we use the last one
        printf("Warning! Add capture to tracker process queue failed! Frame: %d\n", frame_count);

        snprintf(buffer, 128, "%s/%04d_masksAK.png", OUTPUT, frame_count);
        cv::imwrite(buffer, last_mask);
        syncFileWriter.writeStamp(frame_count);
      }
      continue;
    }

    // get capture for body tracker
    k4abt_frame_t body_frame = NULL;
    k4a_wait_result_t pop_frame_result = k4abt_tracker_pop_result(tracker, &body_frame, K4A_WAIT_INFINITE);
    if (pop_frame_result == K4A_WAIT_RESULT_TIMEOUT)
    {
      // It should never hit timeout when K4A_WAIT_INFINITE is set.
      printf("Error! Pop body frame result timeout!\n");
      exit(1);
    }
    else if (pop_frame_result != K4A_WAIT_RESULT_SUCCEEDED)
    {
      printf("Pop body frame result failed!\n");
      exit(1);
    }

    // go to next iteration if we are not writing
    if (!writing)
    {
      k4abt_frame_release(body_frame);
      continue;
    }

    // Process and write mask
    k4a_image_t body_index_map = k4abt_frame_get_body_index_map(body_frame);
    k4a_capture_t input_capture = k4abt_frame_get_capture(body_frame);
    k4a_image_t depth_image = k4a_capture_get_depth_image(input_capture);

    // transform body index into mask
    Mat depth_mask = index_to_opencv(body_index_map);
    cv::bitwise_not(depth_mask, depth_mask);

    VERIFY(k4a_image_create_from_buffer(
               K4A_IMAGE_FORMAT_CUSTOM8,
               depth_image_width_pixels,
               depth_image_height_pixels,
               depth_image_width_pixels * (int)sizeof(uint8_t),
               depth_mask.data,
               (size_t)(depth_mask.total() * depth_mask.elemSize()),
               NULL,
               NULL,
               &mask_image),
           "Could not create image from buffer!");

    VERIFY(k4a_transformation_depth_image_to_color_camera_custom(
               transformation,
               depth_image,
               mask_image,
               transformed_depth_image,
               transformed_mask_image,
               K4A_TRANSFORMATION_INTERPOLATION_TYPE_NEAREST,
               0),
           "Coud not transform depth image!");

    // Get mask, depth and color images
    Mat mask = index_to_opencv(transformed_mask_image);

    Mat rs_mask;
    cv::resize(mask, rs_mask, cv::Size(1920, 1080));
    snprintf(buffer, 128, "%s/%04d_masksAK.png", OUTPUT, frame_count);
    last_mask = rs_mask;
    cv::imwrite(buffer, rs_mask);
    syncFileWriter.writeStamp(frame_count);

    //imshow("BodyTracking", mask);
    //waitKey(1);

    k4a_image_release(mask_image);
    k4a_image_release(depth_image);
    k4a_image_release(body_index_map);
    k4a_capture_release(input_capture);

    k4abt_frame_release(body_frame);
  }

  printf("Finished body tracking processing!\n");

  // Clean up body tracker
  k4abt_tracker_shutdown(tracker);
  k4abt_tracker_destroy(tracker);

  k4a_playback_close(playback_handle);

  syncFileWriter.close();
  return 0;
}