Transforming LDR denoiser to HDR (C)

This example demonstrates how to transform LDR denosier to HDR denoiser using tone-mapping as preprocessing and gamma-correction as postprocessing.

#include "rml/RadeonML.h"
#include "rml/RadeonML_graph.h"

#include <inttypes.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define RML_CHECK(STATUS)                                 \
        do                                                \
        {                                                 \
                if (STATUS != RML_OK)                     \
                {                                         \
                        printf("%s\n", rmlGetLastError());\
                        exit(EXIT_FAILURE);               \
                }                                         \
        } while (0)

#define CHECK(STATUS)                       \
        do                                  \
        {                                   \
                if (!(STATUS))              \
                {                           \
                        exit(EXIT_FAILURE); \
                }                           \
        } while (0)

#define NUM_INPUTS 4
#define NHWC_RANK 4

/*
 * Create operation that holds information aboud input data
 *
 * @param graph - graph where operation is created
 * @param name - unique operation name
 * @param shape - shape of input tensor
 * @param row_index - index in array of shape
 * @return - created placeholder operation
 */
rml_op CreatePlaceholderOp(rml_graph graph, const char* name, const uint32_t* shape)
{
        // Create placeholder operation
        rml_op_desc input_desc = {
                RML_OP_PLACEHOLDER, name, .placeholder = {RML_DTYPE_FLOAT32, RML_LAYOUT_NHWC}};
        memcpy(input_desc.placeholder.tensor_info.shape, shape, NHWC_RANK * sizeof(int));
        rml_op op_placeholder = NULL;
        RML_CHECK(rmlCreateOperation(graph, &input_desc, &op_placeholder));
        return op_placeholder;
}

/*
 * Create operation that stores scalar int data
 *
 * @param graph - graph where operation is created
 * @param name - unique operation name
 * @param dtype - tensor data type
 * @param value - scalar value
 * @return - created scalar operation
 */

rml_op CreateScalarOp(rml_graph graph, const char* name, rml_dtype dtype, const void* value)
{
        rml_op_desc op_desc = {RML_OP_CONST, name, .constant = {{dtype, RML_LAYOUT_SCALAR}, value}};
        rml_op op_const = NULL;
        RML_CHECK(rmlCreateOperation(graph, &op_desc, &op_const));
        return op_const;
}
/*
 * Create specilalized unary operation
 *
 * @param graph - graph where operation is created
 * @param name - unique operation name
 * @param op_type - type of unary operation
 * @param input - input operation
 * @return - created unary operation
 */
rml_op CreateUnaryOp(rml_graph graph, const char* name, rml_op_type op_type, rml_op input)
{
        // Create unary operation
        rml_op_desc op_desc = {op_type, name, .unary = {input}};
        rml_op op_unary = NULL;
        RML_CHECK(rmlCreateOperation(graph, &op_desc, &op_unary));
        return op_unary;
}

/*
 * Create specilalized binary operation
 *
 * @param graph - graph where operation is created
 * @param name - unique operation name
 * @param op_type - type of binary operation
 * @param input1 - input1 operation
 * @param input2 - input2 operation
 * @return - created binary operation
 */
rml_op CreateBinaryOp(rml_graph graph,
                                          const char* name,
                                          rml_op_type op_type,
                                          rml_op input1,
                                          rml_op input2)
{
        // Create binary operation
        rml_op_desc op_desc = {op_type, name, .binary = {input1, input2}};
        rml_op op_binary = NULL;
        RML_CHECK(rmlCreateOperation(graph, &op_desc, &op_binary));
        return op_binary;
}

/*
 * Create preprocessing graph and connect it with base graph
 *
 * @param graph - base graph to be connected with preprocesssing graph
 * @param input_names - list of unique input names of preprocessing graph
 * @param input_shapes - list of input shapes of preprocessing graph
 * @return - connected graph
 */
rml_graph ConnectPreprocessingGraph(const rml_graph graph,
                                                                        const char* input_names[],
                                                                        const rml_tensor_info* input_infos)
{
        // Preprocessing graph includes exponential tone-mapping
        // ldr_color = beta - exp(alpha * hdr_color)
        // alpha = -1.0
        // beta = 1.0
        /* Create a context */
        rml_graph preprocess_graph = NULL;
        RML_CHECK(rmlCreateGraph(&preprocess_graph));

        // Create placeholder per input
        rml_op color_op = CreatePlaceholderOp(preprocess_graph, input_names[0], input_infos[0].shape);
        rml_op albedo_op = CreatePlaceholderOp(preprocess_graph, input_names[1], input_infos[1].shape);
        rml_op depth_op = CreatePlaceholderOp(preprocess_graph, input_names[2], input_infos[2].shape);
        rml_op normal_op = CreatePlaceholderOp(preprocess_graph, input_names[3], input_infos[3].shape);

        // Create alpha
        float alpha_const = -1.0f;
        rml_op alpha_op = CreateScalarOp(preprocess_graph, "alpha", RML_DTYPE_FLOAT32, &alpha_const);

        // Create beta
        float beta_const = 1.0f;
        rml_op beta_op = CreateScalarOp(preprocess_graph, "beta", RML_DTYPE_FLOAT32, &beta_const);

        // Create mul, exp and sub opearions
        rml_op mul_op = CreateBinaryOp(preprocess_graph, "mul", RML_OP_MUL, alpha_op, color_op);
        rml_op exp_op = CreateUnaryOp(preprocess_graph, "exp", RML_OP_EXP, mul_op);
        rml_op sub_op = CreateBinaryOp(preprocess_graph, "sub", RML_OP_SUB, beta_op, exp_op);

        // Create axis for concatenation
        int axis = -1;
        rml_op axis_op = CreateScalarOp(preprocess_graph, "concat/axis", RML_DTYPE_INT32, &axis);

        // Create inputs for concatenation
        rml_op inputs[NUM_INPUTS] = {sub_op, albedo_op, depth_op, normal_op};
        // Concatenate tone-mapped color with albedo, depth and normal
        rml_op_desc concat_desc = {RML_OP_CONCAT, "concat", .concat = {NUM_INPUTS, inputs, axis_op}};

        rml_op op_concat = NULL;
        RML_CHECK(rmlCreateOperation(preprocess_graph, &concat_desc, &op_concat));

        // Get tail graph inputs
        rml_strings tail_inputs;
        RML_CHECK(rmlGetGraphInputNames(graph, &tail_inputs));

        // Get head graph outputs
        rml_strings head_outputs;
        RML_CHECK(rmlGetGraphOutputNames(preprocess_graph, &head_outputs));

        // Connect preprocessing graph with base graph
        rml_graph connected_graph = NULL;
        RML_CHECK(rmlConnectGraphs(preprocess_graph,
                                                           graph,
                                                           1,
                                                           &head_outputs.items[0],
                                                           &tail_inputs.items[0],
                                                           &connected_graph));

        return connected_graph;
}

/*
 * Create postprocessing graph and connect it with base graph
 *
 * @param graph - base graph to be connected with postprocessing graph
 * @param input_name - unique input name of postprocessing graph
 * @param input_shape - input shape of postprocessing graph
 * @return connected graph
 */
rml_graph ConnectPostprocessingGraph(rml_graph graph,
                                                                         const char* input_name,
                                                                         const uint32_t* input_shape)
{
        // Postprocessing graph includes gamma-correction
        // ldr_color = (clip(ldr_color, 0, 1)) ^ (gamma)
        // gamma = 0.4
        rml_graph postprocess_graph = NULL;
        RML_CHECK(rmlCreateGraph(&postprocess_graph));

        // Create placeholder for color
        rml_op input_op = CreatePlaceholderOp(postprocess_graph, input_name, input_shape);

        // Clip color
        rml_op_desc clip_desc = {RML_OP_CLIP, "clip", .clip = {input_op, 0.f, 1.f}};

        rml_op clip_op = NULL;
        RML_CHECK(rmlCreateOperation(postprocess_graph, &clip_desc, &clip_op));

        // Create gamma
        float gamma_const = 0.4f;
        rml_op gamma_op = CreateScalarOp(postprocess_graph, "gamma", RML_DTYPE_FLOAT32, &gamma_const);

        // Create pow operation
        rml_op pow_op = CreateBinaryOp(postprocess_graph, "pow", RML_OP_POW, clip_op, gamma_op);

        // Get tail graph inputs
        rml_strings tail_inputs;
        RML_CHECK(rmlGetGraphInputNames(postprocess_graph, &tail_inputs));

        // Get head graph outputs
        rml_strings head_outputs;
        RML_CHECK(rmlGetGraphOutputNames(graph, &head_outputs));

        // Connect base graph with postprocessing graph
        rml_graph connected_graph = NULL;
        RML_CHECK(rmlConnectGraphs(graph,
                                                           postprocess_graph,
                                                           1,
                                                           &head_outputs.items[0],
                                                           &tail_inputs.items[0],
                                                           &connected_graph));
        return connected_graph;
}

/*
 * Read input from file
 * Must be free memory after caller of this function
 *
 * @param input_file - name of input file
 * @return - string content of file
 */
void* ReadInput(const char* input_file)
{
        void* buffer;
        FILE* file = fopen(input_file, "rb");
        CHECK(file != NULL);
        printf("Reading data from file: %s\n", input_file);

        fseek(file, 0, SEEK_END);
        long length = ftell(file);
        fseek(file, 0, SEEK_SET);
        buffer = malloc((length) * sizeof(char));
        CHECK(buffer != NULL);
        size_t num_read = fread(buffer, sizeof(char), length, file);
        CHECK(num_read == length);
        printf("Input data size: %zu\n", num_read);

        fclose(file);
        return buffer;
}

/*
 * Write output to file
 *
 * @param output_file - name of the output file
 * @param output - output data
 * @param count - number of element in output
 */
void WriteOutput(const char* output_file, const void* output, const size_t count)
{
        FILE* file = fopen(output_file, "wb");
        CHECK(file != NULL);
        printf("Writing result to file: %s\n", output_file);

        size_t count_written = fwrite(output, sizeof(char), count, file);
        CHECK(count_written == count);
        printf("Output data size: %zu\n", count_written);

        fclose(file);
}

/*
 * This sample demonstrates how ldr-denosier could be converted to hdr-denoiser
 * using tone-mapping as preprocessing and gamma-correction as postprocessing
 */
int main()
{
        /* Set model parameters */
#if defined(_WIN32)
        const rml_char* model_path = L"path/model";
#else
        const rml_char* model_path = "path/model";
#endif

        // Set input files
        const char* input_files[] = {
                "path/color",
                "path/albedo",
                "path/depth",
                "path/normal",
        };

        // Set output file
        const char* output_file = "path/output";

        // Set input names
        const char* input_names[] = {"hdr-color", "albedo", "depth", "normal"};

        // Set input shapes
        const rml_tensor_info input_infos[NUM_INPUTS] = {
                {RML_DTYPE_FLOAT32, RML_LAYOUT_NHWC, {1, 600, 800, 3}},
                {RML_DTYPE_FLOAT32, RML_LAYOUT_NHWC, {1, 600, 800, 3}},
                {RML_DTYPE_FLOAT32, RML_LAYOUT_NHWC, {1, 600, 800, 1}},
                {RML_DTYPE_FLOAT32, RML_LAYOUT_NHWC, {1, 600, 800, 2}}};

        // Create a context
        rml_context context = NULL;
        RML_CHECK(rmlCreateDefaultContext(NULL, &context));

        // Load model as a mutable graph
        // model input - 9-channel 800x600 image (3-channel ldr-color,
        //                                        3-channel albedo,
        //                                        1-channel depth,
        //                                        2-channel normal)
        // model output - 3-channel 800x600 ldr image
        rml_graph graph = NULL;
        RML_CHECK(rmlLoadGraphFromFile(model_path, &graph));

        // Add preprocessing of base model inputs
        // Before we can use ldr-denoiser for hdr-data, we should adjust hdr-color
        // using tone-mapping and concatenate it with albedo, depth and normal
        rml_graph graph_with_preprocess = ConnectPreprocessingGraph(graph, input_names, input_infos);
        rmlReleaseGraph(graph);

        // Add postprocessing of base model outputs
        // We should also apply gamma-correction for denoised image
        rml_graph full_graph =
                ConnectPostprocessingGraph(graph_with_preprocess, "input", input_infos[0].shape);
        rmlReleaseGraph(graph_with_preprocess);

        // Create immutable model from connected graphs
        rml_model model = NULL;
        RML_CHECK(rmlCreateModelFromGraph(context, full_graph, &model));

        /* Release the graph */
        rmlReleaseGraph(full_graph);

        // Set up input info
        size_t i;
        for (i = 0; i < NUM_INPUTS; i++)
        {
                RML_CHECK(rmlSetModelInputInfo(model, input_names[i], &input_infos[i]));
        }

        // Check memory info
        rml_memory_info memory_info;
        RML_CHECK(rmlGetModelMemoryInfo(model, &memory_info));

        // Create and fill the input tensors
        rml_tensor inputs[NUM_INPUTS];
        for (i = 0; i < NUM_INPUTS; i++)
        {
                rml_tensor input;
                RML_CHECK(rmlCreateTensor(context, &input_infos[i], RML_ACCESS_MODE_WRITE_ONLY, &input));
                size_t data_size = 0;
                void* data = NULL;
                RML_CHECK(rmlMapTensor(input, &data, &data_size));
                void* file_data = ReadInput(input_files[i]);
                memcpy(data, file_data, data_size);
                free(file_data);
                RML_CHECK(rmlUnmapTensor(input, data));
                inputs[i] = input;
        }

        // Set model inputs
        for (i = 0; i < NUM_INPUTS; i++)
        {
                RML_CHECK(rmlSetModelInput(model, input_names[i], inputs[i]));
        }

        // Get output tensor information
        rml_tensor_info output_info;
        RML_CHECK(rmlGetModelOutputInfo(model, NULL, &output_info));

        // Create the output tensor
        rml_tensor output_tensor = NULL;
        RML_CHECK(rmlCreateTensor(context, &output_info, RML_ACCESS_MODE_READ_ONLY, &output_tensor));

        // Set model output
        RML_CHECK(rmlSetModelOutput(model, NULL, output_tensor));

        // Run the inference
        RML_CHECK(rmlInfer(model));

        // Get data from output tensor
        size_t output_size;
        void* output_data = NULL;
        RML_CHECK(rmlMapTensor(output_tensor, &output_data, &output_size));

        // Unmap output data
        RML_CHECK(rmlUnmapTensor(output_tensor, &output_data));

        // Write the output
        WriteOutput(output_file, output_data, output_size);

        /* Release the input and output tensors */
        rmlReleaseTensor(output_tensor);
        for (i = 0; i < NUM_INPUTS; i++)
        {
                rmlReleaseTensor(inputs[i]);
        }

        /* Release the model */
        rmlReleaseModel(model);

        /* Release the context */
        rmlReleaseContext(context);

        return 0;
}