ImageFormat: refactor to incremental write II

Grok Compression · Grok Compression · commit 057ffcaff6f5 · 2026-04-20T16:40:50.000-04:00
diff --git a/src/lib/codec/apps/GrkDecompress.cpp b/src/lib/codec/apps/GrkDecompress.cpp
@@ -108,7 +108,9 @@ static void decompress_help_display(void)
           grk_version());
   fputs(decompress_help_text, stdout);
 }
-GrkDecompress::GrkDecompress() : storeToDisk(true), imageFormat(nullptr), messenger_(nullptr) {}
+GrkDecompress::GrkDecompress()
+    : storeToDisk(true), incrementalWriteActive_(false), imageFormat(nullptr), messenger_(nullptr)
+{}
 GrkDecompress::~GrkDecompress(void)
 {
   delete imageFormat;
@@ -921,6 +923,15 @@ static bool grkWriteStripCallback(uint32_t workerId, grk_io_buf buffer, void* us
   return imageFormat->writeStrip(workerId, buffer);
 }
 
+static bool grkWriteBandCallback(uint32_t yBegin, uint32_t yEnd, grk_image* image, void* user_data)
+{
+  if(!user_data || !image)
+    return false;
+  auto imageFormat = (IImageFormat*)user_data;
+  imageFormat->setImage(image);
+  return imageFormat->writeImageBand(yBegin, yEnd);
+}
+
 bool GrkDecompress::writeHeader(grk_plugin_decompress_callback_info* info)
 {
   if(!storeToDisk)
@@ -1106,8 +1117,21 @@ int GrkDecompress::preProcess(grk_plugin_decompress_callback_info* info)
   // 3a. initialize writer before decompress so it's ready for incremental output
   if(!writeInit(info))
     goto cleanup;
-  if(!writeHeader(info))
-    goto cleanup;
+  // enable incremental band writing when postProcess is a no-op
+  incrementalWriteActive_ = false;
+  if(storeToDisk && !parameters->single_tile_decompress && !info->init_decompressors_func &&
+     grk_image_is_post_process_no_op(info->image))
+  {
+    // write header early — metadata is already final since postProcess is a no-op
+    if(!writeHeader(info))
+      goto cleanup;
+    auto fmt = info->format_private ? (IImageFormat*)info->format_private : imageFormat;
+    if(fmt->supportsIncrementalBandWrite())
+    {
+      grk_decompress_set_band_callback(info->codec, grkWriteBandCallback, fmt);
+      incrementalWriteActive_ = true;
+    }
+  }
   // 3b. decompress one particular tile
   if(parameters->single_tile_decompress)
   {
@@ -1280,10 +1304,22 @@ int GrkDecompress::postProcess(grk_plugin_decompress_callback_info* info)
       }
     }
     auto outfileStr = outfile ? std::string(outfile) : "";
-    if(!fmt->writeImage())
+    if(incrementalWriteActive_)
     {
-      spdlog::error("Outfile {} not generated", outfileStr);
-      goto cleanup;
+      // bands already written during decompress; just finalize
+    }
+    else
+    {
+      if(!(fmt->getWriteState() & IMAGE_FORMAT_HEADER_WRITTEN) && !fmt->writeHeader())
+      {
+        spdlog::error("Encode header failed for {}", outfileStr);
+        goto cleanup;
+      }
+      if(!fmt->writeImage())
+      {
+        spdlog::error("Outfile {} not generated", outfileStr);
+        goto cleanup;
+      }
     }
     if(!fmt->writeFinish())
     {
diff --git a/src/lib/codec/apps/GrkDecompress.h b/src/lib/codec/apps/GrkDecompress.h
@@ -80,6 +80,7 @@ class GrkDecompress
   int shmBatchDecompress(DecompressInitParams* initParams);
 
   bool storeToDisk;
+  bool incrementalWriteActive_;
   IImageFormat* imageFormat;
   grk_plugin::Messenger* messenger_;
 };
diff --git a/src/lib/codec/formats/IImageFormat.h b/src/lib/codec/formats/IImageFormat.h
@@ -87,6 +87,14 @@ class IImageWriter
   virtual bool writeInit(grk_image* image, const std::string& filename, uint32_t compression_level,
                          uint32_t concurrency) = 0;
 
+  /**
+   * @brief Replace the image pointer used by subsequent write calls.
+   *
+   * Used for incremental band writing when the source image differs
+   * from the one originally passed to writeInit().
+   */
+  virtual void setImage(grk_image* image) = 0;
+
   /**
    * @brief Writes format-specific header
    *
@@ -117,6 +125,11 @@ class IImageWriter
    */
   virtual bool writeImageBand(uint32_t yBegin, uint32_t yEnd) = 0;
 
+  /**
+   * @brief Returns true if this format supports incremental band writing via writeImageBand().
+   */
+  virtual bool supportsIncrementalBandWrite(void) const = 0;
+
   /**
    * @brief Writes a single strip of pre-packed pixel data (push-based, incremental).
    *
diff --git a/src/lib/codec/formats/ImageFormat.cpp b/src/lib/codec/formats/ImageFormat.cpp
@@ -78,6 +78,14 @@ bool ImageFormat::writeInit(grk_image* image, const std::string& filename,
 
   return true;
 }
+void ImageFormat::setImage(grk_image* image)
+{
+  image_ = image;
+}
+bool ImageFormat::supportsIncrementalBandWrite(void) const
+{
+  return false;
+}
 /***
  * library-orchestrated pixel encoding
  */
diff --git a/src/lib/codec/formats/ImageFormat.h b/src/lib/codec/formats/ImageFormat.h
@@ -40,6 +40,8 @@ class ImageFormat : public IImageFormat
   void reclaim(uint32_t workerId, grk_io_buf pixels);
   virtual bool writeInit(grk_image* image, const std::string& filename, uint32_t compression_level,
                          uint32_t concurrency) override;
+  virtual void setImage(grk_image* image) override;
+  bool supportsIncrementalBandWrite(void) const override;
   virtual bool writeStrip(uint32_t workerId, grk_io_buf pixels) override;
   virtual bool writeFinish(void) override;
   uint32_t getWriteState(void) override;
diff --git a/src/lib/codec/formats/JPEGFormat.h b/src/lib/codec/formats/JPEGFormat.h
@@ -95,6 +95,10 @@ class JPEGFormat : public ImageFormat
   bool writeHeader(void) override;
   bool writeImage() override;
   bool writeImageBand(uint32_t yBegin, uint32_t yEnd) override;
+  bool supportsIncrementalBandWrite(void) const override
+  {
+    return true;
+  }
   using ImageFormat::writeStrip;
   bool writeFinish(void) override;
   grk_image* readImage(const std::string& filename, grk_cparameters* parameters) override;
diff --git a/src/lib/codec/formats/PNGFormat.h b/src/lib/codec/formats/PNGFormat.h
@@ -43,6 +43,10 @@ class PNGFormat : public ImageFormat
   bool writeHeader(void) override;
   bool writeImage() override;
   bool writeImageBand(uint32_t yBegin, uint32_t yEnd) override;
+  bool supportsIncrementalBandWrite(void) const override
+  {
+    return true;
+  }
   using ImageFormat::writeStrip;
   bool writeFinish(void) override;
   grk_image* readImage(const std::string& filename, grk_cparameters* parameters) override;
@@ -129,11 +133,6 @@ bool PNGFormat<T>::writeHeader(void)
       break;
     if(image_->comps[0].sgnd != image_->comps[i].sgnd)
       break;
-    if(!image_->comps[i].data)
-    {
-      spdlog::error("imagetopng: component {} is null.", i);
-      return false;
-    }
   }
   if(i != nr_comp_)
   {
diff --git a/src/lib/codec/formats/PNMFormat.h b/src/lib/codec/formats/PNMFormat.h
@@ -52,6 +52,13 @@ class PNMFormat : public ImageFormat
   bool writeHeader(void) override;
   bool writeImage() override;
   bool writeImageBand(uint32_t yBegin, uint32_t yEnd) override;
+  bool supportsIncrementalBandWrite(void) const override
+  {
+    // Split PGM files (forceSplit + single component) need the full image at once
+    if(forceSplit && image_ && image_->decompress_num_comps <= 1)
+      return false;
+    return true;
+  }
   using ImageFormat::writeStrip;
   bool writeFinish(void) override;
   grk_image* readImage(const std::string& filename, grk_cparameters* parameters) override;
diff --git a/src/lib/codec/formats/TIFFFormat.h b/src/lib/codec/formats/TIFFFormat.h
@@ -639,6 +639,10 @@ class TIFFFormat : public ImageFormat
    */
   bool writeImage() override;
   bool writeImageBand(uint32_t yBegin, uint32_t yEnd) override;
+  bool supportsIncrementalBandWrite(void) const override
+  {
+    return true;
+  }
   /***
    * library-orchestrated pixel encoding
    */
diff --git a/src/lib/core/codestream/IDecompressor.h b/src/lib/core/codestream/IDecompressor.h
@@ -63,6 +63,17 @@ struct IDecompressor
    */
   virtual void init(grk_decompress_parameters* param) = 0;
 
+  /**
+   * @brief Sets a band-completion callback for incremental writing.
+   *
+   * Must be called after readHeader() but before decompress().
+   */
+  virtual void setBandCallback(grk_io_band_callback callback, void* user_data)
+  {
+    (void)callback;
+    (void)user_data;
+  }
+
   /**
    * @brief Gets the @ref grk_progression_state for a tile
    *
diff --git a/src/lib/core/codestream/decompress/CodeStreamDecompress.cpp b/src/lib/core/codestream/decompress/CodeStreamDecompress.cpp
@@ -140,6 +140,8 @@ void CodeStreamDecompress::init(grk_decompress_parameters* parameters)
   ioBufferCallback_ = core->io_buffer_callback;
   ioUserData_ = core->io_user_data;
   grkRegisterReclaimCallback_ = core->io_register_client_callback;
+  ioBandCallback_ = core->io_band_callback;
+  ioBandUserData_ = core->io_band_user_data;
 
   // Initialize compressed chunk cache for LRU + network fetches
   if((core->tile_cache_strategy & GRK_TILE_CACHE_LRU) && stream_->getFetcher())
@@ -151,6 +153,11 @@ void CodeStreamDecompress::init(grk_decompress_parameters* parameters)
 
   postReadHeader();
 }
+void CodeStreamDecompress::setBandCallback(grk_io_band_callback callback, void* user_data)
+{
+  ioBandCallback_ = callback;
+  ioBandUserData_ = user_data;
+}
 
 // Multi Tile //////////////////////////////////////////////////////////
 
@@ -213,6 +220,53 @@ bool CodeStreamDecompress::decompress(grk_plugin_tile* tile)
 
   multiTileComposite_->postReadHeader(&cp_);
   tileCache_->init(cp_.t_grid_width_ * cp_.t_grid_height_);
+  // create TileCompletion for band callback if set after header read
+  if(ioBandCallback_ && !tileCompletion_)
+  {
+    // Use unreduced image bounds so that TileCompletion computes the same tile grid
+    // (t_grid_width_ x t_grid_height_) as the codec.  Tile indices from
+    // tileProcessor->getIndex() are based on the unreduced grid.
+    Rect32 unreducedBounds(cp_.tx0_, cp_.ty0_,
+                           cp_.tx0_ + cp_.t_grid_width_ * cp_.t_width_,
+                           cp_.ty0_ + cp_.t_grid_height_ * cp_.t_height_);
+    auto slatedRect = tilesToDecompress_.getSlatedTileRect();
+    nextBandTileY_ = slatedRect.y0;
+    pendingBands_.clear();
+    uint8_t reduce = cp_.codingParams_.dec_.reduce_;
+    // Compute reduced-resolution pixel extents for Y clamping.
+    // multiTileComposite_ bounds are unreduced at this point; apply reduction.
+    auto compositeBounds = multiTileComposite_->getBounds();
+    uint32_t regionY0 = ceildivpow2<uint32_t>(compositeBounds.y0, reduce);
+    uint32_t regionY1 = ceildivpow2<uint32_t>(compositeBounds.y1, reduce);
+    uint32_t reducedTileHeight = ceildivpow2<uint32_t>(cp_.t_height_, reduce);
+    uint32_t reducedTy0 = ceildivpow2<uint32_t>(cp_.ty0_, reduce);
+    tileCompletion_ = std::make_unique<TileCompletion>(
+        tileCache_.get(), unreducedBounds, cp_.t_width_, cp_.t_height_,
+        [this, regionY0, regionY1, reducedTileHeight, reducedTy0](uint16_t tileIndexBegin,
+                                                                   uint16_t) {
+          uint16_t numTileCols = tileCompletion_->getNumTileCols();
+          uint16_t tileY = tileIndexBegin / numTileCols;
+          uint32_t tileGlobalYBegin = reducedTy0 + (uint32_t)tileY * reducedTileHeight;
+          uint32_t tileGlobalYEnd =
+              std::min(reducedTy0 + ((uint32_t)tileY + 1) * reducedTileHeight, regionY1);
+          uint32_t yBegin = std::max(tileGlobalYBegin, regionY0) - regionY0;
+          uint32_t yEnd = std::min(tileGlobalYEnd, regionY1) - regionY0;
+          if(yEnd <= yBegin)
+            return;
+          // buffer this band and drain in order
+          std::lock_guard<std::mutex> lock(bandOrderMutex_);
+          pendingBands_[tileY] = {yBegin, yEnd};
+          while(pendingBands_.count(nextBandTileY_))
+          {
+            auto& band = pendingBands_[nextBandTileY_];
+            if(!ioBandCallback_(band.yBegin, band.yEnd, scratchImage_.get(), ioBandUserData_))
+              success_ = false;
+            pendingBands_.erase(nextBandTileY_);
+            nextBandTileY_++;
+          }
+        },
+        nullptr, tilesToDecompress_.getSlatedTileRect());
+  }
   auto slatedTiles = tilesToDecompress_.getSlatedTiles();
   if(!decompressImpl(slatedTiles))
     return false;
@@ -968,8 +1022,16 @@ std::function<void()> CodeStreamDecompress::postMultiTile(void)
     }
     if(!cp_.codingParams_.dec_.skipAllocateComposite_)
     {
-      scratchImage_->transferDataTo(multiTileComposite_.get());
-      success_ = postProcess(multiTileComposite_.get());
+      if(ioBandCallback_)
+      {
+        // incremental band writes already consumed the data from scratchImage_;
+        // skip transfer and postProcess
+      }
+      else
+      {
+        scratchImage_->transferDataTo(multiTileComposite_.get());
+        success_ = postProcess(multiTileComposite_.get());
+      }
     }
   };
 }
diff --git a/src/lib/core/codestream/decompress/CodeStreamDecompress.h b/src/lib/core/codestream/decompress/CodeStreamDecompress.h
@@ -20,6 +20,7 @@
 #include "ConcurrentQueue.h"
 #include "TFSingleton.h"
 #include "CompressedChunkCache.h"
+#include <map>
 
 namespace grk
 {
@@ -53,6 +54,8 @@ class CodeStreamDecompress final : public CodeStream, public IDecompressor
 
   void init(grk_decompress_parameters* param) override;
 
+  void setBandCallback(grk_io_band_callback callback, void* user_data) override;
+
   grk_progression_state getProgressionState(uint16_t tile_index) override;
 
   bool setProgressionState(grk_progression_state state) override;
@@ -401,6 +404,21 @@ class CodeStreamDecompress final : public CodeStream, public IDecompressor
    */
   void* ioUserData_ = nullptr;
 
+  /**
+   * @brief callback invoked when a tile-row band is ready for incremental writing
+   */
+  grk_io_band_callback ioBandCallback_ = nullptr;
+  void* ioBandUserData_ = nullptr;
+
+  // band ordering for incremental writes
+  std::mutex bandOrderMutex_;
+  uint16_t nextBandTileY_ = 0;
+  struct PendingBand_
+  {
+    uint32_t yBegin, yEnd;
+  };
+  std::map<uint16_t, PendingBand_> pendingBands_;
+
   /**
    * @brief callback to reclaim io buffers
    *
diff --git a/src/lib/core/codestream/decompress/CodeStreamDecompress_ReadMarkers.cpp b/src/lib/core/codestream/decompress/CodeStreamDecompress_ReadMarkers.cpp
@@ -229,8 +229,9 @@ bool CodeStreamDecompress::readHeader(grk_header_info* headerInfo)
       headerImage_->num_precision = headerInfo->num_precision;
     }
     headerImage_->copyHeaderTo(multiTileComposite_.get());
-    multiTileComposite_->postReadHeader(&cp_);
     multiTileComposite_->validateColourSpace();
+    headerImage_->validateColourSpace();
+    multiTileComposite_->postReadHeader(&cp_);
     uint32_t num_threads = (uint32_t)TFSingleton::num_threads();
     coderPool_.makeCoders(num_threads, 6, 6, [this]() -> std::shared_ptr<t1::ICoder> {
       return std::shared_ptr<t1::ICoder>(
diff --git a/src/lib/core/fileformat/decompress/FileFormatJP2Decompress.cpp b/src/lib/core/fileformat/decompress/FileFormatJP2Decompress.cpp
@@ -265,6 +265,10 @@ CodingParams* FileFormatJP2Decompress::getCodingParams(void)
 {
   return codeStream->getCodingParams();
 }
+void FileFormatJP2Decompress::setBandCallback(grk_io_band_callback callback, void* user_data)
+{
+  codeStream->setBandCallback(callback, user_data);
+}
 
 bool FileFormatJP2Decompress::read_xml(uint8_t* p_xml_data, uint32_t xml_size)
 {
diff --git a/src/lib/core/fileformat/decompress/FileFormatJP2Decompress.h b/src/lib/core/fileformat/decompress/FileFormatJP2Decompress.h
@@ -63,6 +63,7 @@ class FileFormatJP2Decompress final : public FileFormatJP2Family, public IDecomp
   void wait(grk_wait_swath* swath) override;
   void scheduleSwathCopy(const grk_wait_swath* swath, grk_swath_buffer* buf) override;
   void waitSwathCopy() override;
+  void setBandCallback(grk_io_band_callback callback, void* user_data) override;
   CodingParams* getCodingParams(void);
 
 private:
diff --git a/src/lib/core/fileformat/decompress/FileFormatMJ2Decompress.cpp b/src/lib/core/fileformat/decompress/FileFormatMJ2Decompress.cpp
@@ -692,6 +692,14 @@ bool FileFormatMJ2Decompress::decompressTile(uint16_t tile_index)
 void FileFormatMJ2Decompress::dump([[maybe_unused]] uint32_t flag,
                                    [[maybe_unused]] FILE* outputFileStream) {};
 void FileFormatMJ2Decompress::wait([[maybe_unused]] grk_wait_swath* swath) {}
+void FileFormatMJ2Decompress::setBandCallback(grk_io_band_callback callback, void* user_data)
+{
+  for(auto& sc : sampleCodeStreams_)
+  {
+    if(sc.codeStream)
+      sc.codeStream->setBandCallback(callback, user_data);
+  }
+}
 bool FileFormatMJ2Decompress::readHeader(grk_header_info* header_info)
 {
   bool rc = FileFormatJP2Family::readHeader(header_info, headerImage_);
diff --git a/src/lib/core/fileformat/decompress/FileFormatMJ2Decompress.h b/src/lib/core/fileformat/decompress/FileFormatMJ2Decompress.h
@@ -41,6 +41,8 @@ class FileFormatMJ2Decompress : public IDecompressor, public FileFormatMJ2
   void dump(uint32_t flag, FILE* outputFileStream) override;
   void wait(grk_wait_swath* swath) override;
 
+  void setBandCallback(grk_io_band_callback callback, void* user_data) override;
+
   uint32_t getNumSamples(void) override;
   bool decompressSample(uint32_t sampleIndex) override;
   GrkImage* getSampleImage(uint32_t sampleIndex) override;
diff --git a/src/lib/core/grok.cpp b/src/lib/core/grok.cpp
diff --git a/src/lib/core/grok.h b/src/lib/core/grok.h
diff --git a/src/lib/core/util/GrkImage.cpp b/src/lib/core/util/GrkImage.cpp
diff --git a/src/lib/core/util/GrkImage.h b/src/lib/core/util/GrkImage.h
