DatasetExporter.cpp

#include "template.h"

#ifdef ML_DATASET

#ifdef _WIN32
// Manually define the specific constants TinyEXR needs
#ifndef CP_UTF8
#define CP_UTF8 65001
#endif

// Forward declare the Windows function TinyEXR is looking for
extern "C" __declspec(dllimport) int __stdcall MultiByteToWideChar(
    unsigned int CodePage,
    unsigned long dwFlags,
    const char* lpMultiByteStr,
    int cbMultiByte,
    wchar_t* lpWideCharStr,
    int cchWideChar);
#endif

#define TINYEXR_USE_MINIZ 0
#include "zlib.h"
#define TINYEXR_IMPLEMENTATION
#include "tinyexr.h"

#include "DatasetExporter.h"
#include <vector>
#include <string>
#include <iostream>

bool DatasetExporter::SaveNoisyFrame(const std::string& filename, const FrameData& data)
{
    EXRHeader header;
    InitEXRHeader(&header);

    EXRImage image;
    InitEXRImage(&image);

    // Total channels: Distance(1), Illum_RGB(3), Normal_XYZ(3) = 7
    constexpr int num_channels = 7;
    image.num_channels = num_channels;
    image.width = data.width;
    image.height = data.height;

    // 1. MUST BE STRICTLY ALPHABETICAL for mrv2 and other viewers to map correctly
    const char* names[] = {
        "Distance",       // 1
        "Illumination_B", // 2
        "Illumination_G", // 3
        "Illumination_R", // 4
        "Nx",             // 5
        "Ny",             // 6
        "Nz",             // 7
    };

    // 2. Prepare pixel arrays
    std::vector<float> channels[num_channels];
    for (int i = 0; i < num_channels; ++i) channels[i].resize(data.width * data.height);

    const float scale = 1.f / (data.spp - 1);

    for (int i = 0; i < data.width * data.height; ++i)
    {
        float4 rgb = RGB8_to_RGBAF32(data.rgb[i]);

        channels[0][i] = data.distance[i] * scale;    // Distance

        // Illumination RGB (indices 4, 5, 6)
        channels[1][i] = data.illumination[i * 3 + 2] * scale; // Illum B
        channels[2][i] = data.illumination[i * 3 + 1] * scale; // Illum G
        channels[3][i] = data.illumination[i * 3 + 0] * scale; // Illum R

        // Normals (indices 7, 8, 9)
        channels[4][i] = data.normal[i * 3 + 0]; // Nx
        channels[5][i] = data.normal[i * 3 + 1]; // Ny
        channels[6][i] = data.normal[i * 3 + 2]; // Nz
    }

    // 3. Link pointers to the EXR structure
    float* image_ptr[num_channels];
    EXRChannelInfo channel_infos[num_channels];
    for (int i = 0; i < num_channels; ++i)
    {
        image_ptr[i] = channels[i].data();
        strncpy(channel_infos[i].name, names[i], 255);
        channel_infos[i].name[255] = '\0';
    }

    image.images = (unsigned char**)image_ptr;
    header.num_channels = num_channels;
    header.channels = channel_infos;

    // 4. Set pixel types
    header.pixel_types = new int[num_channels];
    header.requested_pixel_types = new int[num_channels];
    for (int i = 0; i < num_channels; ++i)
    {
        header.pixel_types[i] = TINYEXR_PIXELTYPE_FLOAT;
        header.requested_pixel_types[i] = TINYEXR_PIXELTYPE_FLOAT;
    }

    // 5. Save to disk
    const char* err = nullptr;
    int ret = SaveEXRImageToFile(&image, &header, filename.c_str(), &err);

    // Cleanup
    delete[] header.pixel_types;
    delete[] header.requested_pixel_types;

    if (ret != TINYEXR_SUCCESS)
    {
        std::cerr << "EXR Error: " << err << std::endl;
        FreeEXRErrorMessage(err);
        return false;
    }
    return true;
}

bool DatasetExporter::SaveHDFrame(const std::string& filename, const FrameData& data)
{
    EXRHeader header;
    InitEXRHeader(&header);

    EXRImage image;
    InitEXRImage(&image);

    // Total channels: Illum_RGB(3) = 3
    constexpr int num_channels = 3;
    image.num_channels = num_channels;
    image.width = data.width;
    image.height = data.height;

    // 1. MUST BE STRICTLY ALPHABETICAL for mrv2 and other viewers to map correctly
    const char* names[] = {
        "Illumination_B", // 1
        "Illumination_G", // 2
        "Illumination_R", // 3
    };

    // 2. Prepare pixel arrays
    std::vector<float> channels[num_channels];
    for (int i = 0; i < num_channels; ++i) channels[i].resize(data.width * data.height);
    
    const float scale = 1.f / (data.spp - 1);

    for (int i = 0; i < data.width * data.height; ++i)
    {
        // Illumination RGB (indices 4, 5, 6)
        channels[0][i] = data.illumination[i * 3 + 2] * scale; // Illum B
        channels[1][i] = data.illumination[i * 3 + 1] * scale; // Illum G
        channels[2][i] = data.illumination[i * 3 + 0] * scale; // Illum R
    }

    // 3. Link pointers to the EXR structure
    float* image_ptr[num_channels];
    EXRChannelInfo channel_infos[num_channels];
    for (int i = 0; i < num_channels; ++i)
    {
        image_ptr[i] = channels[i].data();
        strncpy(channel_infos[i].name, names[i], 255);
        channel_infos[i].name[255] = '\0';
    }

    image.images = (unsigned char**)image_ptr;
    header.num_channels = num_channels;
    header.channels = channel_infos;

    // 4. Set pixel types
    header.pixel_types = new int[num_channels];
    header.requested_pixel_types = new int[num_channels];
    for (int i = 0; i < num_channels; ++i)
    {
        header.pixel_types[i] = TINYEXR_PIXELTYPE_FLOAT;
        header.requested_pixel_types[i] = TINYEXR_PIXELTYPE_FLOAT;
    }

    // 5. Save to disk
    const char* err = nullptr;
    int ret = SaveEXRImageToFile(&image, &header, filename.c_str(), &err);

    // Cleanup
    delete[] header.pixel_types;
    delete[] header.requested_pixel_types;

    if (ret != TINYEXR_SUCCESS)
    {
        std::cerr << "EXR Error: " << err << std::endl;
        FreeEXRErrorMessage(err);
        return false;
    }
    return true;
}

void DatasetExporter::PushNoisy(const FrameData& data)
{
    m_lowq = data;
    std::string filename = m_path + std::to_string(data.frame) + "_lq_.exr";
    SaveNoisyFrame(filename, data);
}

void DatasetExporter::PushHD(const FrameData& data)
{
    m_highq = data;
    std::string filename = m_path + std::to_string(data.frame) + "_hq_.exr";
    SaveHDFrame(filename, data);
}

#endif // ML_DATASET