304 lines
13 KiB
C++
304 lines
13 KiB
C++
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||
|
// Licensed under the MIT License.
|
||
|
|
||
|
#include <assert.h>
|
||
|
#include <iostream>
|
||
|
|
||
|
#include <k4a/k4a.h>
|
||
|
#include <k4abt.h>
|
||
|
|
||
|
#define VERIFY(result, error) \
|
||
|
if(result != K4A_RESULT_SUCCEEDED) \
|
||
|
{ \
|
||
|
printf("%s \n - (File: %s, Function: %s, Line: %d)\n", error, __FILE__, __FUNCTION__, __LINE__); \
|
||
|
exit(1); \
|
||
|
} \
|
||
|
|
||
|
void print_body_index_map_middle_line(k4a_image_t body_index_map)
|
||
|
{
|
||
|
uint8_t* body_index_map_buffer = k4a_image_get_buffer(body_index_map);
|
||
|
|
||
|
// Given body_index_map pixel type should be uint8, the stride_byte should be the same as width
|
||
|
// TODO: Since there is no API to query the byte-per-pixel information, we have to compare the width and stride to
|
||
|
// know the information. We should replace this assert with proper byte-per-pixel query once the API is provided by
|
||
|
// K4A SDK.
|
||
|
assert(k4a_image_get_stride_bytes(body_index_map) == k4a_image_get_width_pixels(body_index_map));
|
||
|
|
||
|
int middle_line_num = k4a_image_get_height_pixels(body_index_map) / 2;
|
||
|
body_index_map_buffer = body_index_map_buffer + middle_line_num * k4a_image_get_width_pixels(body_index_map);
|
||
|
|
||
|
printf("BodyIndexMap at Line %d:\n", middle_line_num);
|
||
|
for (int i = 0; i < k4a_image_get_width_pixels(body_index_map); i++)
|
||
|
{
|
||
|
printf("%u, ", *body_index_map_buffer);
|
||
|
body_index_map_buffer++;
|
||
|
}
|
||
|
printf("\n");
|
||
|
}
|
||
|
|
||
|
// Transform skeleton results from 3d depth space to 2d color image space
|
||
|
inline bool transform_joint_from_depth_3d_to_color_2d(
|
||
|
const k4a_calibration_t* calibration,
|
||
|
k4a_float3_t joint_in_depth_space,
|
||
|
k4a_float2_t& joint_in_color_2d)
|
||
|
{
|
||
|
int valid;
|
||
|
VERIFY(k4a_calibration_3d_to_2d(
|
||
|
calibration,
|
||
|
&joint_in_depth_space,
|
||
|
K4A_CALIBRATION_TYPE_DEPTH,
|
||
|
K4A_CALIBRATION_TYPE_COLOR,
|
||
|
&joint_in_color_2d,
|
||
|
&valid), "Failed to project 3d joint from depth space to 2d color image space!");
|
||
|
|
||
|
return valid != 0;
|
||
|
}
|
||
|
|
||
|
// Transform body index map results from depth space to color space
|
||
|
void transform_body_index_map_from_depth_to_color(
|
||
|
k4a_transformation_t transformation_handle,
|
||
|
const k4a_image_t depth_image,
|
||
|
const k4a_image_t body_index_map_in_depth_space,
|
||
|
k4a_image_t depth_image_in_color_space,
|
||
|
k4a_image_t body_index_map_in_color_space)
|
||
|
{
|
||
|
// Note:
|
||
|
// 1. Depth image - In order to transform the body index map to color space, the corresponding depth image is
|
||
|
// required to help perform this transformation in 3d.
|
||
|
// 2. Interpolation type - Each pixel value for the body index map represents the body index. It is not interpolatable.
|
||
|
// The interpolation method has to be set to K4A_TRANSFORMATION_INTERPOLATION_TYPE_NEAREST.
|
||
|
// 3. Invalid custom value - Because there is disparity between the depth camera and color camera. There might be
|
||
|
// invalid values during the transform. We want this invalid value to be set to K4ABT_BODY_INDEX_MAP_BACKGROUND.
|
||
|
|
||
|
VERIFY(k4a_transformation_depth_image_to_color_camera_custom(
|
||
|
transformation_handle,
|
||
|
depth_image,
|
||
|
body_index_map_in_depth_space,
|
||
|
depth_image_in_color_space,
|
||
|
body_index_map_in_color_space,
|
||
|
K4A_TRANSFORMATION_INTERPOLATION_TYPE_NEAREST,
|
||
|
K4ABT_BODY_INDEX_MAP_BACKGROUND), "Failed to transform body index map to color space!");
|
||
|
}
|
||
|
|
||
|
bool ProcessArguments(k4abt_tracker_configuration_t& tracker_config, int argc, char** argv)
|
||
|
{
|
||
|
#ifdef _WIN32
|
||
|
printf("Usage: k4abt_camera_space_transform_sample PROCESSING_MODE[CUDA, CPU, DirectML ( default ), or TensorRT](optional) -model MODEL_FILEPATH(optional).\n");
|
||
|
#else
|
||
|
printf("Usage: k4abt_camera_space_transform_sample PROCESSING_MODE[CUDA ( default ), CPU, or TensorRT](optional) -model MODEL_FILEPATH(optional).\n");
|
||
|
#endif
|
||
|
|
||
|
for (int i = 1; i < argc; i++)
|
||
|
{
|
||
|
if (0 == strcmp(argv[i], "TensorRT"))
|
||
|
{
|
||
|
tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_TENSORRT;
|
||
|
}
|
||
|
else if (0 == strcmp(argv[i], "CUDA"))
|
||
|
{
|
||
|
tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_CUDA;
|
||
|
}
|
||
|
else if (0 == strcmp(argv[i], "CPU"))
|
||
|
{
|
||
|
tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_CPU;
|
||
|
}
|
||
|
#ifdef _WIN32
|
||
|
else if (0 == strcmp(argv[i], "DirectML"))
|
||
|
{
|
||
|
tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_DIRECTML;
|
||
|
}
|
||
|
#endif
|
||
|
else if (0 == strcmp(argv[i], "-model"))
|
||
|
{
|
||
|
if( i < argc - 1 )
|
||
|
tracker_config.model_path = argv[++i];
|
||
|
else
|
||
|
{
|
||
|
printf("Error: model filepath missing\n");
|
||
|
return false;
|
||
|
}
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
#ifdef _WIN32
|
||
|
printf("Invalid processing mode ! Accepted values are CUDA, CPU, DirectML ( default ), or TensorRT.\n");
|
||
|
#else
|
||
|
printf("Invalid processing mode ! Accepted values are CUDA ( default ), CPU, or TensorRT.\n");
|
||
|
#endif
|
||
|
return false;
|
||
|
}
|
||
|
}
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
int main(int argc, char** argv)
|
||
|
{
|
||
|
k4a_device_configuration_t device_config = K4A_DEVICE_CONFIG_INIT_DISABLE_ALL;
|
||
|
device_config.depth_mode = K4A_DEPTH_MODE_NFOV_UNBINNED;
|
||
|
device_config.color_resolution = K4A_COLOR_RESOLUTION_720P;
|
||
|
|
||
|
k4a_device_t device;
|
||
|
VERIFY(k4a_device_open(0, &device), "Open K4A Device failed!");
|
||
|
VERIFY(k4a_device_start_cameras(device, &device_config), "Start K4A cameras failed!");
|
||
|
|
||
|
// Make sure to pass in the correct device config for both depth camera and color camera to get the correct sensor calibration
|
||
|
k4a_calibration_t sensor_calibration;
|
||
|
VERIFY(k4a_device_get_calibration(device, device_config.depth_mode, device_config.color_resolution, &sensor_calibration),
|
||
|
"Get depth camera calibration failed!");
|
||
|
|
||
|
// Create transformation handle to perform the body index map space transform
|
||
|
k4a_transformation_t transformation = NULL;
|
||
|
transformation = k4a_transformation_create(&sensor_calibration);
|
||
|
if (transformation == NULL)
|
||
|
{
|
||
|
printf("Failed to create transformation from sensor calibration!");
|
||
|
exit(1);
|
||
|
}
|
||
|
|
||
|
k4abt_tracker_t tracker = NULL;
|
||
|
k4abt_tracker_configuration_t tracker_config = K4ABT_TRACKER_CONFIG_DEFAULT;
|
||
|
if (!ProcessArguments(tracker_config, argc, argv))
|
||
|
exit(1);
|
||
|
VERIFY(k4abt_tracker_create(&sensor_calibration, tracker_config, &tracker), "Body tracker initialization failed!");
|
||
|
|
||
|
// Preallocated the buffers to hold the depth image in color space and the body index map in color space
|
||
|
int color_image_width_pixels = sensor_calibration.color_camera_calibration.resolution_width;
|
||
|
int color_image_height_pixels = sensor_calibration.color_camera_calibration.resolution_height;
|
||
|
|
||
|
k4a_image_t depth_image_in_color_space = NULL;
|
||
|
VERIFY(k4a_image_create(K4A_IMAGE_FORMAT_DEPTH16,
|
||
|
color_image_width_pixels,
|
||
|
color_image_height_pixels,
|
||
|
color_image_width_pixels * (int)sizeof(uint16_t),
|
||
|
&depth_image_in_color_space), "Failed to create empty image for the depth image in color space");
|
||
|
|
||
|
k4a_image_t body_index_map_in_color_space = NULL;
|
||
|
VERIFY(k4a_image_create(K4A_IMAGE_FORMAT_CUSTOM8,
|
||
|
color_image_width_pixels,
|
||
|
color_image_height_pixels,
|
||
|
color_image_width_pixels * (int)sizeof(uint8_t),
|
||
|
&body_index_map_in_color_space), "Failed to create empty image for the body index map in color space");
|
||
|
|
||
|
int frame_count = 0;
|
||
|
do
|
||
|
{
|
||
|
k4a_capture_t sensor_capture;
|
||
|
k4a_wait_result_t get_capture_result = k4a_device_get_capture(device, &sensor_capture, K4A_WAIT_INFINITE);
|
||
|
if (get_capture_result == K4A_WAIT_RESULT_SUCCEEDED)
|
||
|
{
|
||
|
frame_count++;
|
||
|
|
||
|
printf("Start processing frame %d\n", frame_count);
|
||
|
|
||
|
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");
|
||
|
break;
|
||
|
}
|
||
|
else if (queue_capture_result == K4A_WAIT_RESULT_FAILED)
|
||
|
{
|
||
|
printf("Error! Add capture to tracker process queue failed!\n");
|
||
|
break;
|
||
|
}
|
||
|
|
||
|
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_SUCCEEDED)
|
||
|
{
|
||
|
uint32_t num_bodies = k4abt_frame_get_num_bodies(body_frame);
|
||
|
printf("%u bodies are detected!\n", num_bodies);
|
||
|
|
||
|
// Transform each 3d joints from 3d depth space to 2d color image space
|
||
|
for (uint32_t i = 0; i < num_bodies; i++)
|
||
|
{
|
||
|
printf("Person[%u]:\n", i);
|
||
|
k4abt_skeleton_t skeleton;
|
||
|
VERIFY(k4abt_frame_get_body_skeleton(body_frame, i, &skeleton), "Get body from body frame failed!");
|
||
|
for (int joint_id = 0; joint_id < (int)K4ABT_JOINT_COUNT; joint_id++)
|
||
|
{
|
||
|
k4a_float2_t joint_in_color_2d;
|
||
|
bool valid = transform_joint_from_depth_3d_to_color_2d(
|
||
|
&sensor_calibration,
|
||
|
skeleton.joints[joint_id].position,
|
||
|
joint_in_color_2d);
|
||
|
if (valid)
|
||
|
{
|
||
|
printf("Joint[%d]: Pixel Location at Color Image ( %f, %f) \n",
|
||
|
joint_id, joint_in_color_2d.v[0], joint_in_color_2d.v[1]);
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
printf("Joint[%d]: Invalid Pixel Location \n", joint_id);
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Transform the body index map from the depth space to color space
|
||
|
k4a_image_t body_index_map_in_depth_space = k4abt_frame_get_body_index_map(body_frame);
|
||
|
if (body_index_map_in_depth_space != NULL)
|
||
|
{
|
||
|
// Depth image is needed in order to perform the body index map space transform
|
||
|
k4a_image_t depth_image = k4a_capture_get_depth_image(sensor_capture);
|
||
|
|
||
|
transform_body_index_map_from_depth_to_color(
|
||
|
transformation,
|
||
|
depth_image,
|
||
|
body_index_map_in_depth_space,
|
||
|
depth_image_in_color_space,
|
||
|
body_index_map_in_color_space);
|
||
|
|
||
|
print_body_index_map_middle_line(body_index_map_in_color_space);
|
||
|
k4a_image_release(body_index_map_in_depth_space);
|
||
|
k4a_image_release(depth_image);
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
printf("Error: Fail to generate bodyindex map!\n");
|
||
|
}
|
||
|
|
||
|
k4abt_frame_release(body_frame);
|
||
|
}
|
||
|
else 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");
|
||
|
break;
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
printf("Pop body frame result failed!\n");
|
||
|
break;
|
||
|
}
|
||
|
}
|
||
|
else if (get_capture_result == K4A_WAIT_RESULT_TIMEOUT)
|
||
|
{
|
||
|
// It should never hit time out when K4A_WAIT_INFINITE is set.
|
||
|
printf("Error! Get depth frame time out!\n");
|
||
|
break;
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
printf("Get depth capture returned error: %d\n", get_capture_result);
|
||
|
break;
|
||
|
}
|
||
|
|
||
|
} while (frame_count < 100);
|
||
|
|
||
|
printf("Finished body tracking processing!\n");
|
||
|
|
||
|
k4a_image_release(depth_image_in_color_space);
|
||
|
k4a_image_release(body_index_map_in_color_space);
|
||
|
|
||
|
k4a_transformation_destroy(transformation);
|
||
|
k4abt_tracker_shutdown(tracker);
|
||
|
k4abt_tracker_destroy(tracker);
|
||
|
k4a_device_stop_cameras(device);
|
||
|
k4a_device_close(device);
|
||
|
|
||
|
return 0;
|
||
|
}
|