main.cpp

#include <k4a/k4a.h>
#include <turbojpeg.h>
#include <fstream>
#include <string>
#include <vector>
#include "transformation_helpers.h"

static bool load_depth_image(const char *file_path, int width, int height, k4a_image_t *depth_image)
{
    size_t image_size = width * height * sizeof(uint16_t);
    std::ifstream file(file_path, std::ios::binary);
    if (!file)
    {
        printf("Failed to open depth file: %s\n", file_path);
        return false;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_image_create(K4A_IMAGE_FORMAT_DEPTH16, width, height, width * sizeof(uint16_t), depth_image))
    {
        printf("Failed to create depth image\n");
        return false;
    }

    file.read(reinterpret_cast<char *>(k4a_image_get_buffer(*depth_image)), image_size);
    if (!file)
    {
        printf("Failed to read depth data from file\n");
        k4a_image_release(*depth_image);
        return false;
    }

    return true;
}

static bool load_color_image(const char *file_path, k4a_image_t *color_image)
{
    tjhandle tjHandle = tjInitDecompress();
    if (tjHandle == nullptr)
    {
        printf("Failed to initialize TurboJPEG decompressor\n");
        return false;
    }

    // Read JPEG file into memory
    std::ifstream file(file_path, std::ios::binary | std::ios::ate);
    if (!file)
    {
        printf("Failed to open color image file: %s\n", file_path);
        return false;
    }

    size_t file_size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<unsigned char> jpeg_buffer(file_size);
    file.read(reinterpret_cast<char *>(jpeg_buffer.data()), file_size);
    if (!file)
    {
        printf("Failed to read color image file\n");
        return false;
    }

    int width, height, jpegSubsamp;
    if (tjDecompressHeader2(tjHandle, jpeg_buffer.data(), jpeg_buffer.size(), &width, &height, &jpegSubsamp) != 0)
    {
        printf("Failed to read JPEG header: %s\n", tjGetErrorStr());
        tjDestroy(tjHandle);
        return false;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_image_create(K4A_IMAGE_FORMAT_COLOR_BGRA32, width, height, width * 4 * sizeof(uint8_t), color_image))
    {
        printf("Failed to create color image\n");
        tjDestroy(tjHandle);
        return false;
    }

    if (tjDecompress2(tjHandle,
                      jpeg_buffer.data(),
                      jpeg_buffer.size(),
                      k4a_image_get_buffer(*color_image),
                      width,
                      0, // pitch
                      height,
                      TJPF_BGRA,
                      TJFLAG_FASTDCT | TJFLAG_FASTUPSAMPLE) != 0)
    {
        printf("Failed to decompress color image: %s\n", tjGetErrorStr());
        k4a_image_release(*color_image);
        tjDestroy(tjHandle);
        return false;
    }

    tjDestroy(tjHandle);
    return true;
}

static int transform(const char *color_image_path, const char *depth_image_path, std::string output_filename)
{
    int returnCode = 1;
    k4a_transformation_t transformation = NULL;
    k4a_image_t depth_image = NULL;
    k4a_image_t color_image = NULL;
    k4a_image_t transformed_depth_image = NULL;
    k4a_image_t point_cloud_image = NULL;
    k4a_calibration_t calibration;
    k4a_device_t device = NULL;
    k4a_device_configuration_t config = K4A_DEVICE_CONFIG_INIT_DISABLE_ALL;

    // Extract parameters from file names
    std::string color_image_str = std::string(color_image_path);
    std::string depth_image_str = std::string(depth_image_path);

    // Set color resolution based on color image file name
    size_t p_pos = color_image_str.rfind('P');
    if (p_pos == std::string::npos || color_image_str[p_pos + 1] != '.')
    {
        printf("Failed to extract color resolution from file name: %s\n", color_image_path);
        return returnCode;
    }
    std::string resolution_str = color_image_str.substr(p_pos - 4, 4);
    if (resolution_str == "2160")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_2160P;
    }
    else if (resolution_str == "1440")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_1440P;
    }
    else if (resolution_str == "1080")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_1080P;
    }
    else if (resolution_str == "G720")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_720P;
    }
    else if (resolution_str == "3072")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_3072P;
    }
    else if (resolution_str == "1536")
    {
        config.color_resolution = K4A_COLOR_RESOLUTION_1536P;
    }
    else
    {
        printf("Unsupported color resolution: %s\n", resolution_str.c_str());
        return returnCode;
    }

    // Set depth mode based on depth image file name
    size_t underscore_pos = depth_image_str.find('_');
    if (underscore_pos == std::string::npos || underscore_pos < 4)
    {
        printf("Failed to extract depth mode from file name: %s\n", depth_image_path);
        return returnCode;
    }
    std::string depth_mode_code = depth_image_str.substr(underscore_pos - 4, 4);
    std::string depth_mode_str = depth_image_str.substr(underscore_pos - 4);

    if (depth_mode_str.find("NFOV_BINNED") != std::string::npos)
    {
        config.depth_mode = K4A_DEPTH_MODE_NFOV_2X2BINNED;
    }
    else if (depth_mode_str.find("NFOV_UNBINNED") != std::string::npos)
    {
        config.depth_mode = K4A_DEPTH_MODE_NFOV_UNBINNED;
    }
    else if (depth_mode_str.find("WFOV_BINNED") != std::string::npos)
    {
        config.depth_mode = K4A_DEPTH_MODE_WFOV_2X2BINNED;
    }
    else if (depth_mode_str.find("WFOV_UNBINNED") != std::string::npos)
    {
        config.depth_mode = K4A_DEPTH_MODE_WFOV_UNBINNED;
    }
    else
    {
        printf("Unsupported depth mode: %s\n", depth_mode_str.c_str());
        return returnCode;
    }

    // Set up the Kinect device to obtain the calibration data
    uint32_t device_count = k4a_device_get_installed_count();
    if (device_count == 0)
    {
        printf("No K4A devices found\n");
        return returnCode;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_device_open(K4A_DEVICE_DEFAULT, &device))
    {
        printf("Failed to open device\n");
        return returnCode;
    }

    config.color_format = K4A_IMAGE_FORMAT_COLOR_MJPG;
    config.camera_fps = K4A_FRAMES_PER_SECOND_30;
    config.synchronized_images_only = true;

    if (K4A_RESULT_SUCCEEDED != k4a_device_get_calibration(device, config.depth_mode, config.color_resolution, &calibration))
    {
        printf("Failed to get calibration\n");
        k4a_device_close(device);
        return returnCode;
    }

    transformation = k4a_transformation_create(&calibration);
    if (transformation == NULL)
    {
        printf("Failed to create transformation handle\n");
        k4a_device_close(device);
        return returnCode;
    }

    if (!load_depth_image(depth_image_path, calibration.depth_camera_calibration.resolution_width, calibration.depth_camera_calibration.resolution_height, &depth_image))
    {
        printf("Failed to load depth image\n");
        goto Cleanup;
    }

    if (!load_color_image(color_image_path, &color_image))
    {
        printf("Failed to load color image\n");
        goto Cleanup;
    }

    // Transform depth image to match color image resolution
    if (K4A_RESULT_SUCCEEDED != k4a_image_create(K4A_IMAGE_FORMAT_DEPTH16,
                                                 calibration.color_camera_calibration.resolution_width,
                                                 calibration.color_camera_calibration.resolution_height,
                                                 calibration.color_camera_calibration.resolution_width * sizeof(uint16_t),
                                                 &transformed_depth_image))
    {
        printf("Failed to create transformed depth image\n");
        goto Cleanup;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_transformation_depth_image_to_color_camera(transformation,
                                                                              depth_image,
                                                                              transformed_depth_image))
    {
        printf("Failed to transform depth image to match color camera geometry\n");
        goto Cleanup;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_image_create(K4A_IMAGE_FORMAT_CUSTOM,
                                                 calibration.color_camera_calibration.resolution_width,
                                                 calibration.color_camera_calibration.resolution_height,
                                                 calibration.color_camera_calibration.resolution_width * 3 * sizeof(int16_t),
                                                 &point_cloud_image))
    {
        printf("Failed to create point cloud image\n");
        goto Cleanup;
    }

    if (K4A_RESULT_SUCCEEDED != k4a_transformation_depth_image_to_point_cloud(transformation,
                                                                              transformed_depth_image,
                                                                              K4A_CALIBRATION_TYPE_COLOR,
                                                                              point_cloud_image))
    {
        printf("Failed to compute point cloud\n");
        goto Cleanup;
    }

    tranformation_helpers_write_point_cloud(point_cloud_image, color_image, output_filename.c_str());
    returnCode = 0;

Cleanup:
    if (point_cloud_image != NULL)
    {
        k4a_image_release(point_cloud_image);
    }
    if (depth_image != NULL)
    {
        k4a_image_release(depth_image);
    }
    if (transformed_depth_image != NULL)
    {
        k4a_image_release(transformed_depth_image);
    }
    if (color_image != NULL)
    {
        k4a_image_release(color_image);
    }
    if (transformation != NULL)
    {
        k4a_transformation_destroy(transformation);
    }
    if (device != NULL)
    {
        k4a_device_close(device);
    }
    return returnCode;
}

int main(int argc, char **argv)
{
    if (argc < 4)
    {
        printf("Usage: transform <color_image.jpeg> <depth_image.raw> <output.ply>\n");
        return 1;
    }

    return transform(argv[1], argv[2], argv[3]);
}