modified the constructor and explictly handle the case where no data …

…is passed. Signed-off-by: Philippe Leprince <[email protected]>
clusty · Jul 24, 2024 · 7aa4b10 · 7aa4b10
1 parent b064c47
commit 7aa4b10
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 16 deletions.
diff --git a/src/lib/OpenEXR/ImfCompressor.cpp b/src/lib/OpenEXR/ImfCompressor.cpp
@@ -172,7 +172,7 @@ newTileCompressor (
 
         case ZSTD_COMPRESSION:
 
-            return new ZstdCompressor (hdr);
+            return new ZstdCompressor (hdr, tileLineSize, numTileLines);
 
         default: return 0;
     }

diff --git a/src/lib/OpenEXR/ImfZstdCompressor.cpp b/src/lib/OpenEXR/ImfZstdCompressor.cpp
@@ -19,9 +19,21 @@ std::mutex g_mutex;
 
 OPENEXR_IMF_INTERNAL_NAMESPACE_SOURCE_ENTER
 
-ZstdCompressor::ZstdCompressor (const Header& hdr)
-    : Compressor (hdr), _outBuffer ()
-{}
+ZstdCompressor::ZstdCompressor (
+    const Header& hdr, size_t tileLineSize, size_t numTileLines)
+    : Compressor (hdr)
+    , m_tileLineSize (tileLineSize)
+    , m_numTileLines (numTileLines)
+    , _outBuffer ()
+{
+    // we only need this for deep images to compute the sample count table size.
+    m_isTiled = hdr.hasTileDescription ();
+    if (m_isTiled)
+    {
+        auto td = hdr.tileDescription ();
+        m_tileArea = td.xSize * td.ySize;
+    }
+}
 
 int
 ZstdCompressor::numScanLines () const
@@ -34,9 +46,16 @@ int
 ZstdCompressor::compress (
     const char* inPtr, int inSize, int minY, const char*& outPtr)
 {
+    if (inSize == 0)
+    {
+        outPtr = nullptr;
+        return 0;
+    }
+
     auto             totalChannelSize = 0;
     std::vector<int> sizePerChannel;
-    auto             channels = header ().channels ();
+    auto             hdr      = header ();
+    auto             channels = hdr.channels ();
     for (auto ch = channels.begin (); ch != channels.end (); ++ch)
     {
         const auto size = pixelTypeSize (ch.channel ().type);
@@ -48,15 +67,16 @@ ZstdCompressor::compress (
 
     if (numSamples * totalChannelSize != inSize)
     {
-        // We received fewer data than expected. It probably is because we are processing
-        // the sampleCounts for DeepExr
-        sizePerChannel = {4}; // we compress it as an int
+        // std::cout << "Tile: BAD" << std::endl;
+        // We received fewer data than expected. It probably is because we are
+        // processing the sampleCounts for DeepExr
+        sizePerChannel = {sampleCountTableSize ()};
     }
 
     outPtr = (char*) malloc (inSize);
     {
         std::lock_guard<std::mutex> lock (g_mutex);
-        _outBuffer.push_back (raw_ptr ((char*) outPtr, &free));
+        _outBuffer.push_back (data_ptr ((char*) outPtr, &free));
     }
 
     const auto fullSize = exr_compress_zstd (
@@ -74,15 +94,30 @@ int
 ZstdCompressor::uncompress (
     const char* inPtr, int inSize, int minY, const char*& outPtr)
 {
+    if (inSize == 0)
+    {
+        outPtr = nullptr;
+        return 0;
+    }
+
     auto  read  = (const char*) inPtr;
     void* write = nullptr;
     auto  ret   = exr_uncompress_zstd (read, inSize, &write, 0);
     {
         std::lock_guard<std::mutex> lock (g_mutex);
-        _outBuffer.push_back (raw_ptr ((char*) write, &free));
+        _outBuffer.push_back (data_ptr ((char*) write, &free));
     }
     outPtr = (const char*) write;
     return ret;
 }
 
+int
+ZstdCompressor::sampleCountTableSize ()
+{
+    if (m_isTiled) { return m_tileArea * sizeof (int); }
+    const Header& hdr = header ();
+    return ((hdr.dataWindow ().max.x + 1) - hdr.dataWindow ().min.x) *
+           sizeof (int);
+}
+
 OPENEXR_IMF_INTERNAL_NAMESPACE_SOURCE_EXIT
diff --git a/src/lib/OpenEXR/ImfZstdCompressor.h b/src/lib/OpenEXR/ImfZstdCompressor.h
@@ -17,16 +17,23 @@ OPENEXR_IMF_INTERNAL_NAMESPACE_HEADER_ENTER
 class ZstdCompressor : public Compressor
 {
 public:
-    explicit ZstdCompressor (const Header& hdr);
+    explicit ZstdCompressor (
+        const Header& hdr, size_t tileLineSize = 0, size_t numTileLines = 0);
 
 private:
-    using raw_ptr = std::unique_ptr<char, decltype (&free)>;
-    std::vector<raw_ptr> _outBuffer;
-    int                  numScanLines () const override; // max
-    int                  compress (
-                         const char* inPtr, int inSize, int minY, const char*& outPtr) override;
+    using data_ptr = std::unique_ptr<char, decltype (&free)>;
+    std::vector<data_ptr> _outBuffer;
+    int                   numScanLines () const override; // max
+    int                   compress (
+                          const char* inPtr, int inSize, int minY, const char*& outPtr) override;
     int uncompress (
         const char* inPtr, int inSize, int minY, const char*& outPtr) override;
+    int sampleCountTableSize ();
+
+    int  m_tileLineSize;
+    int  m_numTileLines;
+    int  m_tileArea;
+    bool m_isTiled;
 };
 
 OPENEXR_IMF_INTERNAL_NAMESPACE_HEADER_EXIT