Skip to content

Commit

Permalink
Implement streaming file output in cjxl. (libjxl#2949)
Browse files Browse the repository at this point in the history
  • Loading branch information
@szabadka
szabadka authored Nov 18, 2023
1 parent b2cb82f commit f7b82e0
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 23 deletions.
52 changes: 35 additions & 17 deletions lib/extras/enc/jxl.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,29 @@ bool SetupFrame(JxlEncoder* enc, JxlEncoderFrameSettings* settings,
return true;
}

bool ReadCompressedOutput(JxlEncoder* enc, std::vector<uint8_t>* compressed) {
compressed->clear();
compressed->resize(4096);
uint8_t* next_out = compressed->data();
size_t avail_out = compressed->size() - (next_out - compressed->data());
JxlEncoderStatus result = JXL_ENC_NEED_MORE_OUTPUT;
while (result == JXL_ENC_NEED_MORE_OUTPUT) {
result = JxlEncoderProcessOutput(enc, &next_out, &avail_out);
if (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;
}
}
compressed->resize(next_out - compressed->data());
if (result != JXL_ENC_SUCCESS) {
fprintf(stderr, "JxlEncoderProcessOutput failed.\n");
return false;
}
return true;
}

bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
const std::vector<uint8_t>* jpeg_bytes,
std::vector<uint8_t>* compressed) {
Expand All @@ -102,6 +125,13 @@ bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
return false;
}

if (params.HasOutputProcessor() &&
JXL_ENC_SUCCESS !=
JxlEncoderSetOutputProcessor(enc, params.output_processor)) {
fprintf(stderr, "JxlEncoderSetOutputProcessorfailed\n");
return false;
}

auto settings = JxlEncoderFrameSettingsCreate(enc, nullptr);
size_t option_idx = 0;
if (!SetFrameOptions(params.options, 0, &option_idx, settings)) {
Expand Down Expand Up @@ -314,24 +344,12 @@ bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
}
}
JxlEncoderCloseInput(enc);
// Reading compressed output
compressed->clear();
compressed->resize(4096);
uint8_t* next_out = compressed->data();
size_t avail_out = compressed->size() - (next_out - compressed->data());
JxlEncoderStatus result = JXL_ENC_NEED_MORE_OUTPUT;
while (result == JXL_ENC_NEED_MORE_OUTPUT) {
result = JxlEncoderProcessOutput(enc, &next_out, &avail_out);
if (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;
if (params.HasOutputProcessor()) {
if (JXL_ENC_SUCCESS != JxlEncoderFlushInput(enc)) {
fprintf(stderr, "JxlEncoderAddChunkedFrame() failed.\n");
return false;
}
}
compressed->resize(next_out - compressed->data());
if (result != JXL_ENC_SUCCESS) {
fprintf(stderr, "JxlEncoderProcessOutput failed.\n");
} else if (!ReadCompressedOutput(enc, compressed)) {
return false;
}
return true;
Expand Down
6 changes: 6 additions & 0 deletions lib/extras/enc/jxl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ struct JXLCompressParams {
// If runner_opaque is set, the decoder uses this parallel runner.
JxlParallelRunner runner = JxlThreadParallelRunner;
void* runner_opaque = nullptr;
JxlEncoderOutputProcessor output_processor = {};
JxlDebugImageCallback debug_image = nullptr;
void* debug_image_opaque = nullptr;
JxlEncoderStats* stats = nullptr;
Expand All @@ -73,6 +74,11 @@ struct JXLCompressParams {
void AddFloatOption(JxlEncoderFrameSettingId id, float val) {
options.emplace_back(JXLOption(id, val, 0));
}
bool HasOutputProcessor() const {
return (output_processor.get_buffer != nullptr &&
output_processor.release_buffer != nullptr &&
output_processor.set_finalized_position != nullptr);
}
};

bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
Expand Down
78 changes: 72 additions & 6 deletions tools/cjxl_main.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,9 @@ struct CompressArgs {
"Enable streaming processing of the input file "
"(works only for PPM and PGM input files).",
&streaming_input, &SetBooleanTrue, 3);
cmdline->AddOptionFlag('\0', "streaming_output",
"Enable incremental writing of the output file.",
&streaming_output, &SetBooleanTrue, 3);
cmdline->AddOptionFlag('\0', "disable_output",
"No output file will be written (for benchmarking)",
&disable_output, &SetBooleanTrue, 3);
Expand Down Expand Up @@ -465,6 +468,7 @@ struct CompressArgs {
const char* file_out = nullptr;
jxl::Override print_profile = jxl::Override::kDefault;
bool streaming_input = false;
bool streaming_output = false;

// Decoding source image flags
ColorHintsProxy color_hints_proxy;
Expand Down Expand Up @@ -934,6 +938,56 @@ void ProcessFlags(const jxl::extras::Codec codec,
});
}

struct JxlOutputProcessor {
bool SetOutputPath(const std::string& path) {
outfile.reset(new FileWrapper(path, "wb"));
if (!*outfile) {
fprintf(stderr,
"Could not open %s for writing\n"
"Error: %s",
path.c_str(), strerror(errno));
return false;
}
return true;
}

JxlEncoderOutputProcessor GetOutputProcessor() {
return JxlEncoderOutputProcessor{this, GetBuffer, ReleaseBuffer, Seek,
SetFinalizedPosition};
}

static void* GetBuffer(void* opaque, size_t* size) {
JxlOutputProcessor* self = reinterpret_cast<JxlOutputProcessor*>(opaque);
self->output.resize(*size);
return self->output.data();
}

static void ReleaseBuffer(void* opaque, size_t written_bytes) {
JxlOutputProcessor* self = reinterpret_cast<JxlOutputProcessor*>(opaque);
if (*self->outfile && fwrite(self->output.data(), 1, written_bytes,
*self->outfile) != written_bytes) {
JXL_WARNING("Failed to write %" PRIuS " bytes to output", written_bytes);
}
self->output.clear();
}

static void Seek(void* opaque, uint64_t position) {
JxlOutputProcessor* self = reinterpret_cast<JxlOutputProcessor*>(opaque);
if (*self->outfile && fseek(*self->outfile, position, SEEK_SET) != 0) {
JXL_WARNING("Failed to seek output.");
}
}

static void SetFinalizedPosition(void* opaque, uint64_t finalized_position) {
JxlOutputProcessor* self = reinterpret_cast<JxlOutputProcessor*>(opaque);
self->finalized_position = finalized_position;
}

std::vector<uint8_t> output;
size_t finalized_position = 0;
std::unique_ptr<FileWrapper> outfile;
};

} // namespace tools
} // namespace jpegxl

Expand Down Expand Up @@ -1074,31 +1128,43 @@ int main(int argc, char** argv) {
params.runner_opaque = runner.get();

jpegxl::tools::SpeedStats stats;
jpegxl::tools::JxlOutputProcessor output_processor;
if (args.streaming_output) {
if (args.file_out && !args.disable_output &&
!output_processor.SetOutputPath(args.file_out)) {
return EXIT_FAILURE;
}
params.output_processor = output_processor.GetOutputProcessor();
}
std::vector<uint8_t> compressed;
for (size_t num_rep = 0; num_rep < args.num_reps; ++num_rep) {
const double t0 = jxl::Now();
if (!EncodeImageJXL(params, ppf, jpeg_bytes, &compressed)) {
if (!EncodeImageJXL(params, ppf, jpeg_bytes,
args.streaming_output ? nullptr : &compressed)) {
fprintf(stderr, "EncodeImageJXL() failed.\n");
return EXIT_FAILURE;
}
const double t1 = jxl::Now();
stats.NotifyElapsed(t1 - t0);
stats.SetImageSize(ppf.info.xsize, ppf.info.ysize);
}
size_t compressed_size = args.streaming_output
? output_processor.finalized_position
: compressed.size();

if (args.file_out && !args.disable_output) {
if (!args.streaming_output && args.file_out && !args.disable_output) {
if (!jpegxl::tools::WriteFile(args.file_out, compressed)) {
std::cerr << "Could not write jxl file." << std::endl;
return EXIT_FAILURE;
}
}
if (!args.quiet) {
if (compressed.size() < 100000) {
if (compressed_size < 100000) {
cmdline.VerbosePrintf(0, "Compressed to %" PRIuS " bytes ",
compressed.size());
compressed_size);
} else {
cmdline.VerbosePrintf(0, "Compressed to %.1f kB ",
compressed.size() * 0.001);
compressed_size * 0.001);
}
// For lossless jpeg-reconstruction, we don't print some stats, since we
// don't have easy access to the image dimensions.
Expand All @@ -1107,7 +1173,7 @@ int main(int argc, char** argv) {
}
if (!args.lossless_jpeg) {
const double bpp =
static_cast<double>(compressed.size() * jxl::kBitsPerByte) / pixels;
static_cast<double>(compressed_size * jxl::kBitsPerByte) / pixels;
cmdline.VerbosePrintf(0, "(%.3f bpp%s).\n", bpp / ppf.num_frames(),
ppf.num_frames() == 1 ? "" : "/frame");
JXL_CHECK(stats.Print(num_worker_threads));
Expand Down
4 changes: 4 additions & 0 deletions tools/scripts/roundtrip_test.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,10 @@ main() {
"-e 1 --streaming_input" 0.02
roundtrip_test "jxl/flower/flower_small.rgb.depth8.ppm" \
"-e 1 -d 0.0 --streaming_input" 0.0
roundtrip_test "jxl/flower/flower_small.rgb.depth8.ppm" \
"-e 1 --streaming_output" 0.02
roundtrip_test "jxl/flower/flower_small.rgb.depth8.ppm" \
"-e 1 -d 0.0 --streaming_input --streaming_output" 0.0
roundtrip_test "jxl/flower/flower_cropped.jpg" "-e 1" 0.0

roundtrip_lossless_pnm_test "jxl/flower/flower_small.rgb.depth1.ppm"
Expand Down

0 comments on commit f7b82e0

Please sign in to comment.