moved stuff around, added a grayscale test.
This commit is contained in:
349
util/output/aviwriter.hpp
Normal file
349
util/output/aviwriter.hpp
Normal file
@@ -0,0 +1,349 @@
|
||||
#ifndef AVI_WRITER_HPP
|
||||
#define AVI_WRITER_HPP
|
||||
|
||||
#include <vector>
|
||||
#include <fstream>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <algorithm>
|
||||
#include <filesystem>
|
||||
#include <chrono>
|
||||
|
||||
class AVIWriter {
|
||||
private:
|
||||
#pragma pack(push, 1)
|
||||
struct RIFFChunk {
|
||||
uint32_t chunkId;
|
||||
uint32_t chunkSize;
|
||||
uint32_t format;
|
||||
};
|
||||
|
||||
struct AVIListHeader {
|
||||
uint32_t listId;
|
||||
uint32_t listSize;
|
||||
uint32_t listType;
|
||||
};
|
||||
|
||||
struct AVIMainHeader {
|
||||
uint32_t microSecPerFrame;
|
||||
uint32_t maxBytesPerSec;
|
||||
uint32_t paddingGranularity;
|
||||
uint32_t flags;
|
||||
uint32_t totalFrames;
|
||||
uint32_t initialFrames;
|
||||
uint32_t streams;
|
||||
uint32_t suggestedBufferSize;
|
||||
uint32_t width;
|
||||
uint32_t height;
|
||||
uint32_t reserved[4];
|
||||
};
|
||||
|
||||
struct AVIStreamHeader {
|
||||
uint32_t type;
|
||||
uint32_t handler;
|
||||
uint32_t flags;
|
||||
uint16_t priority;
|
||||
uint16_t language;
|
||||
uint32_t initialFrames;
|
||||
uint32_t scale;
|
||||
uint32_t rate;
|
||||
uint32_t start;
|
||||
uint32_t length;
|
||||
uint32_t suggestedBufferSize;
|
||||
uint32_t quality;
|
||||
uint32_t sampleSize;
|
||||
struct {
|
||||
int16_t left;
|
||||
int16_t top;
|
||||
int16_t right;
|
||||
int16_t bottom;
|
||||
} rcFrame;
|
||||
};
|
||||
|
||||
struct BITMAPINFOHEADER {
|
||||
uint32_t size;
|
||||
int32_t width;
|
||||
int32_t height;
|
||||
uint16_t planes;
|
||||
uint16_t bitCount;
|
||||
uint32_t compression;
|
||||
uint32_t sizeImage;
|
||||
int32_t xPelsPerMeter;
|
||||
int32_t yPelsPerMeter;
|
||||
uint32_t clrUsed;
|
||||
uint32_t clrImportant;
|
||||
};
|
||||
|
||||
struct AVIIndexEntry {
|
||||
uint32_t chunkId;
|
||||
uint32_t flags;
|
||||
uint32_t offset;
|
||||
uint32_t size;
|
||||
};
|
||||
#pragma pack(pop)
|
||||
|
||||
static bool createDirectoryIfNeeded(const std::string& filename) {
|
||||
std::filesystem::path filePath(filename);
|
||||
std::filesystem::path directory = filePath.parent_path();
|
||||
|
||||
if (!directory.empty() && !std::filesystem::exists(directory)) {
|
||||
return std::filesystem::create_directories(directory);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static void writeChunk(std::ofstream& file, uint32_t chunkId, const void* data, uint32_t size) {
|
||||
file.write(reinterpret_cast<const char*>(&chunkId), 4);
|
||||
file.write(reinterpret_cast<const char*>(&size), 4);
|
||||
if (data && size > 0) {
|
||||
file.write(reinterpret_cast<const char*>(data), size);
|
||||
}
|
||||
}
|
||||
|
||||
static void writeList(std::ofstream& file, uint32_t listType, const void* data, uint32_t size) {
|
||||
uint32_t listId = 0x5453494C; // 'LIST'
|
||||
file.write(reinterpret_cast<const char*>(&listId), 4);
|
||||
file.write(reinterpret_cast<const char*>(&size), 4);
|
||||
file.write(reinterpret_cast<const char*>(&listType), 4);
|
||||
if (data && size > 4) {
|
||||
file.write(reinterpret_cast<const char*>(data), size - 4);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
static bool saveAVI(const std::string& filename,
|
||||
const std::vector<std::vector<uint8_t>>& frames,
|
||||
int width, int height, float fps = 30.0f) {
|
||||
if (frames.empty() || width <= 0 || height <= 0 || fps <= 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::cout << "1" << "width: " << width <<
|
||||
"height: " << height << "frame count: " << fps << std::endl;
|
||||
|
||||
// Validate frame sizes
|
||||
size_t expectedFrameSize = width * height * 3;
|
||||
for (const auto& frame : frames) {
|
||||
if (frame.size() != expectedFrameSize) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::cout << "2" << std::endl;
|
||||
// Create directory if needed
|
||||
if (!createDirectoryIfNeeded(filename)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::cout << "3" << std::endl;
|
||||
std::ofstream file(filename, std::ios::binary);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint32_t frameCount = static_cast<uint32_t>(frames.size());
|
||||
uint32_t microSecPerFrame = static_cast<uint32_t>(1000000.0f / fps);
|
||||
|
||||
// Calculate padding for each frame (BMP-style row padding)
|
||||
uint32_t rowSize = (width * 3 + 3) & ~3;
|
||||
uint32_t frameSize = rowSize * height;
|
||||
uint32_t totalDataSize = frameCount * frameSize;
|
||||
|
||||
std::cout << "4" << std::endl;
|
||||
// RIFF AVI header
|
||||
RIFFChunk riffHeader;
|
||||
riffHeader.chunkId = 0x46464952; // 'RIFF'
|
||||
riffHeader.format = 0x20495641; // 'AVI '
|
||||
|
||||
// We'll come back and write the size at the end
|
||||
uint32_t riffStartPos = static_cast<uint32_t>(file.tellp());
|
||||
file.write(reinterpret_cast<const char*>(&riffHeader), sizeof(riffHeader));
|
||||
|
||||
// hdrl list
|
||||
uint32_t hdrlListStart = static_cast<uint32_t>(file.tellp());
|
||||
writeList(file, 0x6C726468, nullptr, 0); // 'hdrl' - we'll fill size later
|
||||
|
||||
std::cout << "5" << std::endl;
|
||||
// avih chunk
|
||||
AVIMainHeader mainHeader;
|
||||
mainHeader.microSecPerFrame = microSecPerFrame;
|
||||
mainHeader.maxBytesPerSec = frameSize * static_cast<uint32_t>(fps);
|
||||
mainHeader.paddingGranularity = 0;
|
||||
mainHeader.flags = 0x000010; // HASINDEX flag
|
||||
mainHeader.totalFrames = frameCount;
|
||||
mainHeader.initialFrames = 0;
|
||||
mainHeader.streams = 1;
|
||||
mainHeader.suggestedBufferSize = frameSize;
|
||||
mainHeader.width = width;
|
||||
mainHeader.height = height;
|
||||
mainHeader.reserved[0] = 0;
|
||||
mainHeader.reserved[1] = 0;
|
||||
mainHeader.reserved[2] = 0;
|
||||
mainHeader.reserved[3] = 0;
|
||||
|
||||
writeChunk(file, 0x68697661, &mainHeader, sizeof(mainHeader)); // 'avih'
|
||||
|
||||
std::cout << "6" << std::endl;
|
||||
// strl list
|
||||
uint32_t strlListStart = static_cast<uint32_t>(file.tellp());
|
||||
writeList(file, 0x6C727473, nullptr, 0); // 'strl' - we'll fill size later
|
||||
|
||||
// strh chunk
|
||||
AVIStreamHeader streamHeader;
|
||||
streamHeader.type = 0x73646976; // 'vids'
|
||||
streamHeader.handler = 0x00000000; // Uncompressed
|
||||
streamHeader.flags = 0;
|
||||
streamHeader.priority = 0;
|
||||
streamHeader.language = 0;
|
||||
streamHeader.initialFrames = 0;
|
||||
streamHeader.scale = 1;
|
||||
streamHeader.rate = static_cast<uint32_t>(fps);
|
||||
streamHeader.start = 0;
|
||||
streamHeader.length = frameCount;
|
||||
streamHeader.suggestedBufferSize = frameSize;
|
||||
streamHeader.quality = 0xFFFFFFFF; // Default quality
|
||||
streamHeader.sampleSize = 0;
|
||||
streamHeader.rcFrame.left = 0;
|
||||
streamHeader.rcFrame.top = 0;
|
||||
streamHeader.rcFrame.right = width;
|
||||
streamHeader.rcFrame.bottom = height;
|
||||
|
||||
writeChunk(file, 0x68727473, &streamHeader, sizeof(streamHeader)); // 'strh'
|
||||
|
||||
// strf chunk
|
||||
BITMAPINFOHEADER bitmapInfo;
|
||||
bitmapInfo.size = sizeof(BITMAPINFOHEADER);
|
||||
bitmapInfo.width = width;
|
||||
bitmapInfo.height = height;
|
||||
bitmapInfo.planes = 1;
|
||||
bitmapInfo.bitCount = 24;
|
||||
bitmapInfo.compression = 0; // BI_RGB - uncompressed
|
||||
bitmapInfo.sizeImage = frameSize;
|
||||
bitmapInfo.xPelsPerMeter = 0;
|
||||
bitmapInfo.yPelsPerMeter = 0;
|
||||
bitmapInfo.clrUsed = 0;
|
||||
bitmapInfo.clrImportant = 0;
|
||||
|
||||
writeChunk(file, 0x66727473, &bitmapInfo, sizeof(bitmapInfo)); // 'strf'
|
||||
|
||||
std::cout << "7" << std::endl;
|
||||
// Update strl list size
|
||||
uint32_t strlListEnd = static_cast<uint32_t>(file.tellp());
|
||||
file.seekp(strlListStart + 4);
|
||||
uint32_t strlListSize = strlListEnd - strlListStart - 8;
|
||||
file.write(reinterpret_cast<const char*>(&strlListSize), 4);
|
||||
file.seekp(strlListEnd);
|
||||
|
||||
std::cout << "8" << std::endl;
|
||||
// Update hdrl list size
|
||||
uint32_t hdrlListEnd = static_cast<uint32_t>(file.tellp());
|
||||
file.seekp(hdrlListStart + 4);
|
||||
uint32_t hdrlListSize = hdrlListEnd - hdrlListStart - 8;
|
||||
file.write(reinterpret_cast<const char*>(&hdrlListSize), 4);
|
||||
file.seekp(hdrlListEnd);
|
||||
|
||||
std::cout << "9" << std::endl;
|
||||
// movi list
|
||||
uint32_t moviListStart = static_cast<uint32_t>(file.tellp());
|
||||
writeList(file, 0x69766F6D, nullptr, 0); // 'movi' - we'll fill size later
|
||||
|
||||
std::vector<AVIIndexEntry> indexEntries;
|
||||
indexEntries.reserve(frameCount);
|
||||
|
||||
// Write frames
|
||||
for (uint32_t i = 0; i < frameCount; ++i) {
|
||||
uint32_t frameStart = static_cast<uint32_t>(file.tellp()) - moviListStart - 4;
|
||||
|
||||
std::cout << "10-" << i << std::endl;
|
||||
// Create padded frame data (BMP-style bottom-to-top with padding)
|
||||
std::vector<uint8_t> paddedFrame(frameSize, 0);
|
||||
const auto& frame = frames[i];
|
||||
uint32_t srcRowSize = width * 3;
|
||||
|
||||
for (int y = 0; y < height; ++y) {
|
||||
int srcY = height - 1 - y; // Flip vertically for BMP format
|
||||
const uint8_t* srcRow = frame.data() + (srcY * srcRowSize);
|
||||
uint8_t* dstRow = paddedFrame.data() + (y * rowSize);
|
||||
memcpy(dstRow, srcRow, srcRowSize);
|
||||
// Padding bytes remain zeros
|
||||
}
|
||||
|
||||
std::cout << "11-" << i << std::endl;
|
||||
// Write frame as '00db' chunk
|
||||
writeChunk(file, 0x62643030, paddedFrame.data(), frameSize); // '00db'
|
||||
|
||||
// Add to index
|
||||
AVIIndexEntry entry;
|
||||
entry.chunkId = 0x62643030; // '00db'
|
||||
entry.flags = 0x00000010; // AVIIF_KEYFRAME
|
||||
entry.offset = frameStart;
|
||||
entry.size = frameSize;
|
||||
indexEntries.push_back(entry);
|
||||
}
|
||||
|
||||
std::cout << "12" << std::endl;
|
||||
// Update movi list size
|
||||
uint32_t moviListEnd = static_cast<uint32_t>(file.tellp());
|
||||
file.seekp(moviListStart + 4);
|
||||
uint32_t moviListSize = moviListEnd - moviListStart - 8;
|
||||
file.write(reinterpret_cast<const char*>(&moviListSize), 4);
|
||||
file.seekp(moviListEnd);
|
||||
|
||||
std::cout << "13" << std::endl;
|
||||
// idx1 chunk - index
|
||||
uint32_t idx1Size = static_cast<uint32_t>(indexEntries.size() * sizeof(AVIIndexEntry));
|
||||
writeChunk(file, 0x31786469, indexEntries.data(), idx1Size); // 'idx1'
|
||||
|
||||
// Update RIFF chunk size
|
||||
uint32_t fileEnd = static_cast<uint32_t>(file.tellp());
|
||||
file.seekp(riffStartPos + 4);
|
||||
uint32_t riffSize = fileEnd - riffStartPos - 8;
|
||||
file.write(reinterpret_cast<const char*>(&riffSize), 4);
|
||||
|
||||
std::cout << "14" << std::endl;
|
||||
return true;
|
||||
}
|
||||
|
||||
// Convenience function to save from individual frame files
|
||||
static bool saveAVIFromFrames(const std::string& filename,
|
||||
const std::vector<std::string>& frameFiles,
|
||||
int width, int height,
|
||||
float fps = 30.0f) {
|
||||
std::vector<std::vector<uint8_t>> frames;
|
||||
frames.reserve(frameFiles.size());
|
||||
|
||||
for (const auto& frameFile : frameFiles) {
|
||||
std::ifstream file(frameFile, std::ios::binary);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Read BMP file and extract pixel data
|
||||
file.seekg(0, std::ios::end);
|
||||
size_t fileSize = file.tellg();
|
||||
file.seekg(0, std::ios::beg);
|
||||
|
||||
std::vector<uint8_t> buffer(fileSize);
|
||||
file.read(reinterpret_cast<char*>(buffer.data()), fileSize);
|
||||
|
||||
// Simple BMP parsing - assumes 24-bit uncompressed BMP
|
||||
if (fileSize < 54 || buffer[0] != 'B' || buffer[1] != 'M') {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Extract pixel data offset from BMP header
|
||||
uint32_t dataOffset = *reinterpret_cast<uint32_t*>(&buffer[10]);
|
||||
if (dataOffset >= fileSize) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Read pixel data (BGR format)
|
||||
std::vector<uint8_t> pixelData(buffer.begin() + dataOffset, buffer.end());
|
||||
frames.push_back(pixelData);
|
||||
}
|
||||
|
||||
return saveAVI(filename, frames, width, height, fps);
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
190
util/output/bmpwriter.hpp
Normal file
190
util/output/bmpwriter.hpp
Normal file
@@ -0,0 +1,190 @@
|
||||
#ifndef BMP_WRITER_HPP
|
||||
#define BMP_WRITER_HPP
|
||||
|
||||
#include <vector>
|
||||
#include <fstream>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <algorithm>
|
||||
#include <filesystem>
|
||||
#include "../vectorlogic/vec3.hpp"
|
||||
|
||||
class BMPWriter {
|
||||
private:
|
||||
#pragma pack(push, 1)
|
||||
struct BMPHeader {
|
||||
uint16_t signature = 0x4D42; // "BM"
|
||||
uint32_t fileSize;
|
||||
uint16_t reserved1 = 0;
|
||||
uint16_t reserved2 = 0;
|
||||
uint32_t dataOffset = 54;
|
||||
};
|
||||
|
||||
struct BMPInfoHeader {
|
||||
uint32_t headerSize = 40;
|
||||
int32_t width;
|
||||
int32_t height;
|
||||
uint16_t planes = 1;
|
||||
uint16_t bitsPerPixel = 24;
|
||||
uint32_t compression = 0;
|
||||
uint32_t imageSize;
|
||||
int32_t xPixelsPerMeter = 0;
|
||||
int32_t yPixelsPerMeter = 0;
|
||||
uint32_t colorsUsed = 0;
|
||||
uint32_t importantColors = 0;
|
||||
};
|
||||
#pragma pack(pop)
|
||||
|
||||
// Helper function to create directory if it doesn't exist
|
||||
static bool createDirectoryIfNeeded(const std::string& filename) {
|
||||
std::filesystem::path filePath(filename);
|
||||
std::filesystem::path directory = filePath.parent_path();
|
||||
|
||||
// If there's a directory component and it doesn't exist, create it
|
||||
if (!directory.empty() && !std::filesystem::exists(directory)) {
|
||||
return std::filesystem::create_directories(directory);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
public:
|
||||
// Save a 2D vector of Vec3 (RGB) colors as BMP
|
||||
// Vec3 components: x = red, y = green, z = blue (values in range [0,1])
|
||||
static bool saveBMP(const std::string& filename, const std::vector<std::vector<Vec3>>& pixels) {
|
||||
if (pixels.empty() || pixels[0].empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int height = static_cast<int>(pixels.size());
|
||||
int width = static_cast<int>(pixels[0].size());
|
||||
|
||||
// Validate that all rows have the same width
|
||||
for (const auto& row : pixels) {
|
||||
if (row.size() != width) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return saveBMP(filename, pixels, width, height);
|
||||
}
|
||||
|
||||
// Alternative interface with width/height and flat vector (row-major order)
|
||||
static bool saveBMP(const std::string& filename, const std::vector<Vec3>& pixels, int width, int height) {
|
||||
if (pixels.size() != width * height) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Convert to 2D vector format
|
||||
std::vector<std::vector<Vec3>> pixels2D(height, std::vector<Vec3>(width));
|
||||
for (int y = 0; y < height; ++y) {
|
||||
for (int x = 0; x < width; ++x) {
|
||||
pixels2D[y][x] = pixels[y * width + x];
|
||||
}
|
||||
}
|
||||
|
||||
return saveBMP(filename, pixels2D, width, height);
|
||||
}
|
||||
|
||||
// Save from 1D vector of uint8_t pixels (BGR order: pixels[i]=b, pixels[i+1]=g, pixels[i+2]=r)
|
||||
static bool saveBMP(const std::string& filename, const std::vector<uint8_t>& pixels, int width, int height) {
|
||||
if (pixels.size() != width * height * 3) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Create directory if needed
|
||||
if (!createDirectoryIfNeeded(filename)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
BMPHeader header;
|
||||
BMPInfoHeader infoHeader;
|
||||
|
||||
int rowSize = (width * 3 + 3) & ~3; // 24-bit, padded to 4 bytes
|
||||
int imageSize = rowSize * height;
|
||||
|
||||
header.fileSize = sizeof(BMPHeader) + sizeof(BMPInfoHeader) + imageSize;
|
||||
infoHeader.width = width;
|
||||
infoHeader.height = height;
|
||||
infoHeader.imageSize = imageSize;
|
||||
|
||||
std::ofstream file(filename, std::ios::binary);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
|
||||
file.write(reinterpret_cast<const char*>(&header), sizeof(header));
|
||||
file.write(reinterpret_cast<const char*>(&infoHeader), sizeof(infoHeader));
|
||||
|
||||
// Write pixel data (BMP stores pixels bottom-to-top)
|
||||
std::vector<uint8_t> row(rowSize, 0);
|
||||
for (int y = height - 1; y >= 0; --y) {
|
||||
const uint8_t* srcRow = pixels.data() + (y * width * 3);
|
||||
|
||||
// Copy and rearrange if necessary (input is already in BGR order)
|
||||
for (int x = 0; x < width; ++x) {
|
||||
int srcOffset = x * 3;
|
||||
int dstOffset = x * 3;
|
||||
|
||||
// Input is already BGR: pixels[i]=b, pixels[i+1]=g, pixels[i+2]=r
|
||||
// So we can copy directly
|
||||
row[dstOffset] = srcRow[srcOffset]; // B
|
||||
row[dstOffset + 1] = srcRow[srcOffset + 1]; // G
|
||||
row[dstOffset + 2] = srcRow[srcOffset + 2]; // R
|
||||
}
|
||||
file.write(reinterpret_cast<const char*>(row.data()), rowSize);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
private:
|
||||
static bool saveBMP(const std::string& filename, const std::vector<std::vector<Vec3>>& pixels, int width, int height) {
|
||||
// Create directory if needed
|
||||
if (!createDirectoryIfNeeded(filename)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
BMPHeader header;
|
||||
BMPInfoHeader infoHeader;
|
||||
|
||||
int rowSize = (width * 3 + 3) & ~3; // 24-bit, padded to 4 bytes
|
||||
int imageSize = rowSize * height;
|
||||
|
||||
header.fileSize = sizeof(BMPHeader) + sizeof(BMPInfoHeader) + imageSize;
|
||||
infoHeader.width = width;
|
||||
infoHeader.height = height;
|
||||
infoHeader.imageSize = imageSize;
|
||||
|
||||
std::ofstream file(filename, std::ios::binary);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
|
||||
file.write(reinterpret_cast<const char*>(&header), sizeof(header));
|
||||
file.write(reinterpret_cast<const char*>(&infoHeader), sizeof(infoHeader));
|
||||
|
||||
// Write pixel data (BMP stores pixels bottom-to-top)
|
||||
std::vector<uint8_t> row(rowSize, 0);
|
||||
for (int y = height - 1; y >= 0; --y) {
|
||||
for (int x = 0; x < width; ++x) {
|
||||
const Vec3& color = pixels[y][x];
|
||||
|
||||
// Convert from [0,1] float to [0,255] uint8_t
|
||||
uint8_t r = static_cast<uint8_t>(std::clamp(color.x * 255.0f, 0.0f, 255.0f));
|
||||
uint8_t g = static_cast<uint8_t>(std::clamp(color.y * 255.0f, 0.0f, 255.0f));
|
||||
uint8_t b = static_cast<uint8_t>(std::clamp(color.z * 255.0f, 0.0f, 255.0f));
|
||||
|
||||
// BMP is BGR order
|
||||
int pixelOffset = x * 3;
|
||||
row[pixelOffset] = b;
|
||||
row[pixelOffset + 1] = g;
|
||||
row[pixelOffset + 2] = r;
|
||||
}
|
||||
file.write(reinterpret_cast<const char*>(row.data()), rowSize);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
233
util/output/jxlwriter.hpp
Normal file
233
util/output/jxlwriter.hpp
Normal file
@@ -0,0 +1,233 @@
|
||||
#ifndef JXL_WRITER_HPP
|
||||
#define JXL_WRITER_HPP
|
||||
|
||||
#include <vector>
|
||||
#include <fstream>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <algorithm>
|
||||
#include <filesystem>
|
||||
#include "vec3.hpp"
|
||||
#include "timing_decorator.hpp"
|
||||
#include <jxl/encode.h>
|
||||
#include <jxl/thread_parallel_runner.h>
|
||||
|
||||
class JXLWriter {
|
||||
private:
|
||||
// Helper function to create directory if it doesn't exist
|
||||
static bool createDirectoryIfNeeded(const std::string& filename) {
|
||||
std::filesystem::path filePath(filename);
|
||||
std::filesystem::path directory = filePath.parent_path();
|
||||
|
||||
// If there's a directory component and it doesn't exist, create it
|
||||
if (!directory.empty() && !std::filesystem::exists(directory)) {
|
||||
return std::filesystem::create_directories(directory);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Helper function to convert Vec3 pixels to interleaved RGB data
|
||||
static std::vector<uint8_t> convertToRGB(const std::vector<std::vector<Vec3>>& pixels, int width, int height) {
|
||||
std::vector<uint8_t> rgbData(width * height * 3);
|
||||
|
||||
for (int y = 0; y < height; ++y) {
|
||||
for (int x = 0; x < width; ++x) {
|
||||
const Vec3& color = pixels[y][x];
|
||||
int offset = (y * width + x) * 3;
|
||||
|
||||
rgbData[offset] = static_cast<uint8_t>(std::clamp(color.x * 255.0f, 0.0f, 255.0f)); // R
|
||||
rgbData[offset + 1] = static_cast<uint8_t>(std::clamp(color.y * 255.0f, 0.0f, 255.0f)); // G
|
||||
rgbData[offset + 2] = static_cast<uint8_t>(std::clamp(color.z * 255.0f, 0.0f, 255.0f)); // B
|
||||
}
|
||||
}
|
||||
|
||||
return rgbData;
|
||||
}
|
||||
|
||||
// Helper function to convert flat Vec3 vector to interleaved RGB data
|
||||
static std::vector<uint8_t> convertToRGB(const std::vector<Vec3>& pixels, int width, int height) {
|
||||
std::vector<uint8_t> rgbData(width * height * 3);
|
||||
|
||||
for (int i = 0; i < width * height; ++i) {
|
||||
const Vec3& color = pixels[i];
|
||||
int offset = i * 3;
|
||||
|
||||
rgbData[offset] = static_cast<uint8_t>(std::clamp(color.x * 255.0f, 0.0f, 255.0f)); // R
|
||||
rgbData[offset + 1] = static_cast<uint8_t>(std::clamp(color.y * 255.0f, 0.0f, 255.0f)); // G
|
||||
rgbData[offset + 2] = static_cast<uint8_t>(std::clamp(color.z * 255.0f, 0.0f, 255.0f)); // B
|
||||
}
|
||||
|
||||
return rgbData;
|
||||
}
|
||||
|
||||
public:
|
||||
// Save a 2D vector of Vec3 (RGB) colors as JXL
|
||||
// Vec3 components: x = red, y = green, z = blue (values in range [0,1])
|
||||
static bool saveJXL(const std::string& filename, const std::vector<std::vector<Vec3>>& pixels,
|
||||
float quality = 90.0f, int effort = 7) {
|
||||
if (pixels.empty() || pixels[0].empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int height = static_cast<int>(pixels.size());
|
||||
int width = static_cast<int>(pixels[0].size());
|
||||
|
||||
// Validate that all rows have the same width
|
||||
for (const auto& row : pixels) {
|
||||
if (row.size() != width) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return saveJXL(filename, pixels, width, height, quality, effort);
|
||||
}
|
||||
|
||||
static bool saveJXL(const std::string& filename, const std::vector<Vec3>& pixels,
|
||||
int width, int height, float quality = 90.0f, int effort = 7) {
|
||||
if (pixels.size() != width * height) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Convert to 2D vector format
|
||||
std::vector<std::vector<Vec3>> pixels2D(height, std::vector<Vec3>(width));
|
||||
for (int y = 0; y < height; ++y) {
|
||||
for (int x = 0; x < width; ++x) {
|
||||
pixels2D[y][x] = pixels[y * width + x];
|
||||
}
|
||||
}
|
||||
|
||||
return saveJXL(filename, pixels2D, width, height, quality, effort);
|
||||
}
|
||||
|
||||
// Save from 1D vector of uint8_t pixels (RGB order: pixels[i]=r, pixels[i+1]=g, pixels[i+2]=b)
|
||||
static bool saveJXL(const std::string& filename, const std::vector<uint8_t>& pixels,
|
||||
int width, int height, float quality = 90.0f, int effort = 7) {
|
||||
TIME_FUNCTION;
|
||||
effort = 1; //forced 1 to test speed
|
||||
if (pixels.size() != width * height * 3) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Create directory if needed
|
||||
if (!createDirectoryIfNeeded(filename)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Create encoder
|
||||
JxlEncoder* enc = JxlEncoderCreate(nullptr);
|
||||
if (!enc) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Create thread pool (use nullptr for single-threaded)
|
||||
void* runner = JxlThreadParallelRunnerCreate(nullptr, JxlThreadParallelRunnerDefaultNumWorkerThreads());
|
||||
if (JxlEncoderSetParallelRunner(enc, JxlThreadParallelRunner, runner) != JXL_ENC_SUCCESS) {
|
||||
JxlThreadParallelRunnerDestroy(runner);
|
||||
JxlEncoderDestroy(enc);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Configure basic encoder settings
|
||||
JxlEncoderFrameSettings* options = JxlEncoderFrameSettingsCreate(enc, nullptr);
|
||||
|
||||
// Set quality/distance (distance = 0 is lossless, higher = more lossy)
|
||||
float distance = (100.0f - quality) / 5.0f;
|
||||
if (quality >= 100.0f) {
|
||||
JxlEncoderSetFrameLossless(options, JXL_TRUE);
|
||||
} else {
|
||||
JxlEncoderSetFrameDistance(options, distance);
|
||||
}
|
||||
|
||||
JxlEncoderFrameSettingsSetOption(options, JXL_ENC_FRAME_SETTING_EFFORT, effort);
|
||||
|
||||
// Set up basic image info
|
||||
JxlBasicInfo basic_info;
|
||||
JxlEncoderInitBasicInfo(&basic_info);
|
||||
basic_info.xsize = width;
|
||||
basic_info.ysize = height;
|
||||
basic_info.bits_per_sample = 8;
|
||||
basic_info.exponent_bits_per_sample = 0;
|
||||
basic_info.uses_original_profile = JXL_FALSE;
|
||||
//basic_info.uses_original_profile = JXL_TRUE;
|
||||
|
||||
if (JxlEncoderSetBasicInfo(enc, &basic_info) != JXL_ENC_SUCCESS) {
|
||||
JxlThreadParallelRunnerDestroy(runner);
|
||||
JxlEncoderDestroy(enc);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Set color encoding to sRGB
|
||||
JxlColorEncoding color_encoding = {};
|
||||
JxlColorEncodingSetToSRGB(&color_encoding, JXL_FALSE); // is_gray = false
|
||||
if (JxlEncoderSetColorEncoding(enc, &color_encoding) != JXL_ENC_SUCCESS) {
|
||||
JxlThreadParallelRunnerDestroy(runner);
|
||||
JxlEncoderDestroy(enc);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Set up pixel format
|
||||
JxlPixelFormat pixel_format = {3, JXL_TYPE_UINT8, JXL_LITTLE_ENDIAN, 0};
|
||||
|
||||
// Add image frame
|
||||
if (JxlEncoderAddImageFrame(options, &pixel_format,
|
||||
(void*)pixels.data(),
|
||||
pixels.size()) != JXL_ENC_SUCCESS) {
|
||||
JxlThreadParallelRunnerDestroy(runner);
|
||||
JxlEncoderDestroy(enc);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Mark the end of input
|
||||
JxlEncoderCloseInput(enc);
|
||||
|
||||
// Compress to buffer
|
||||
std::vector<uint8_t> compressed;
|
||||
// compressed.resize(4096); // Start with 4KB
|
||||
size_t initial_size = width * height / 4;
|
||||
compressed.resize(initial_size);
|
||||
uint8_t* next_out = compressed.data();
|
||||
size_t avail_out = compressed.size();
|
||||
JxlEncoderStatus process_result = JXL_ENC_NEED_MORE_OUTPUT;
|
||||
|
||||
while (process_result == JXL_ENC_NEED_MORE_OUTPUT) {
|
||||
process_result = JxlEncoderProcessOutput(enc, &next_out, &avail_out);
|
||||
if (process_result == JXL_ENC_NEED_MORE_OUTPUT) {
|
||||
size_t offset = next_out - compressed.data();
|
||||
compressed.resize(compressed.size() * 2);
|
||||
next_out = compressed.data() + offset;
|
||||
avail_out = compressed.size() - offset;
|
||||
}
|
||||
}
|
||||
|
||||
bool success = false;
|
||||
if (process_result == JXL_ENC_SUCCESS) {
|
||||
// Write to file
|
||||
size_t compressed_size = next_out - compressed.data();
|
||||
std::ofstream file(filename, std::ios::binary);
|
||||
if (file) {
|
||||
file.write(reinterpret_cast<const char*>(compressed.data()), compressed_size);
|
||||
success = file.good();
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
JxlThreadParallelRunnerDestroy(runner);
|
||||
JxlEncoderDestroy(enc);
|
||||
|
||||
return success;
|
||||
}
|
||||
|
||||
private:
|
||||
static bool saveJXL(const std::string& filename, const std::vector<std::vector<Vec3>>& pixels,
|
||||
int width, int height, float quality, int effort) {
|
||||
if (!createDirectoryIfNeeded(filename)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> rgbData = convertToRGB(pixels, width, height);
|
||||
|
||||
return saveJXL(filename, rgbData, width, height, quality, effort);
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user