Zstd compressor (#3)

Co-authored-by: Philippe Leprince <[email protected]> Signed-off-by: Vlad Lazar <[email protected]> * better multy-type compression * Version the Stream --------- Signed-off-by: Vlad Lazar <[email protected]> Co-authored-by: Philippe Leprince <[email protected]>
AcademySoftwareFoundation · Dec 21, 2023 · 0689612 · 0689612
1 parent fca936a
commit 0689612
Show file tree

Hide file tree

Showing 15 changed files with 321 additions and 39 deletions.
diff --git a/.gitignore b/.gitignore
@@ -57,3 +57,7 @@ docs/_test_images/
 
 # Ignore Bazel generated files
 bazel-* 
+
+cmake-build-debug/
+
+.idea/
diff --git a/src/bin/exrheader/main.cpp b/src/bin/exrheader/main.cpp
@@ -65,6 +65,8 @@ printCompression (Compression c)
 
         case DWAB_COMPRESSION: cout << "dwa, medium scanline blocks"; break;
 
+        case ZSTD_COMPRESSION: cout << "zstd"; break;
+
         default: cout << int (c); break;
     }
 }

diff --git a/src/examples/deepExamples.cpp b/src/examples/deepExamples.cpp
@@ -50,7 +50,7 @@ readDeepScanlineFile (
     //    - allocate the memory requred to store the samples
     //    - read the pixels from the file
     //
-    
+
     DeepScanLineInputFile file (filename);
 
     const Header& header = file.header ();
@@ -118,20 +118,18 @@ readDeepScanlineFile (
     }
 }
 
-unsigned int getPixelSampleCount (int i, int j)
+unsigned int
+getPixelSampleCount (int i, int j)
 {
     // Dummy code creating deep data from a flat image
     return 1;
 }
 
 Array2D<float> testDataZ;
-Array2D<half> testDataA;
+Array2D<half>  testDataA;
 
-void getPixelSampleData(
-    int i,
-    int j,
-    Array2D<float*>& dataZ,
-    Array2D<half*>& dataA)
+void
+getPixelSampleData (int i, int j, Array2D<float*>& dataZ, Array2D<half*>& dataA)
 {
     // Dummy code creating deep data from a flat image
     dataZ[i][j][0] = testDataZ[i][j];
@@ -147,7 +145,8 @@ writeDeepScanlineFile (
 
     Array2D<half*>& dataA,
 
-    Array2D<unsigned int>& sampleCount)
+    Array2D<unsigned int>& sampleCount,
+    Compression            compression = Compression::ZIPS_COMPRESSION)
 
 {
     //
@@ -160,7 +159,7 @@ writeDeepScanlineFile (
     //    - describe the memory layout of the A and Z pixels
     //    - store the pixels in the file
     //
-    
+
     int height = dataWindow.max.y - dataWindow.min.y + 1;
     int width  = dataWindow.max.x - dataWindow.min.x + 1;
 
@@ -169,7 +168,7 @@ writeDeepScanlineFile (
     header.channels ().insert ("Z", Channel (FLOAT));
     header.channels ().insert ("A", Channel (HALF));
     header.setType (DEEPSCANLINE);
-    header.compression () = ZIPS_COMPRESSION;
+    header.compression () = compression;
 
     DeepScanLineOutputFile file (filename, header);
 
@@ -211,7 +210,7 @@ writeDeepScanlineFile (
             dataZ[i][j]       = new float[sampleCount[i][j]];
             dataA[i][j]       = new half[sampleCount[i][j]];
             // Generate data for dataZ and dataA.
-            getPixelSampleData(i, j, dataZ, dataA);
+            getPixelSampleData (i, j, dataZ, dataA);
         }
 
         file.writePixels (1);
@@ -227,30 +226,56 @@ writeDeepScanlineFile (
     }
 }
 
-
-void deepExamples()
+void
+deepExamples ()
 {
     int w = 800;
     int h = 600;
-    
+
     Box2i window;
-    window.min.setValue(0, 0);
-    window.max.setValue(w - 1, h - 1);
-    
-    Array2D<float *> dataZ;
-    dataZ.resizeErase(h, w);
-    
-    Array2D<half *> dataA;
-    dataA.resizeErase(h, w);
-    
+    window.min.setValue (0, 0);
+    window.max.setValue (w - 1, h - 1);
+
+    Array2D<float*> dataZ;
+    dataZ.resizeErase (h, w);
+
+    Array2D<half*> dataA;
+    dataA.resizeErase (h, w);
+
     Array2D<unsigned int> sampleCount;
-    sampleCount.resizeErase(h, w);
-    
+    sampleCount.resizeErase (h, w);
+
     // Create an image to be used as a source for deep data
-    testDataA.resizeErase(h, w);
-    testDataZ.resizeErase(h, w);
-    drawImage2(testDataA, testDataZ, w, h);
-
-    writeDeepScanlineFile("test.deep.exr", window, window, dataZ, dataA, sampleCount);
-    readDeepScanlineFile ("test.deep.exr", window, window, dataZ, dataA, sampleCount);
+    testDataA.resizeErase (h, w);
+    testDataZ.resizeErase (h, w);
+    drawImage2 (testDataA, testDataZ, w, h);
+
+    {
+        writeDeepScanlineFile (
+            "test.deep.exr",
+            window,
+            window,
+            dataZ,
+            dataA,
+            sampleCount,
+            Compression::ZSTD_COMPRESSION);
+    }
+    {
+        writeDeepScanlineFile (
+            "test.zips.exr",
+            window,
+            window,
+            dataZ,
+            dataA,
+            sampleCount,
+            Compression::ZIPS_COMPRESSION);
+    }
+    {
+        readDeepScanlineFile (
+            "test.deep.exr", window, window, dataZ, dataA, sampleCount);
+    }
+    {
+        readDeepScanlineFile (
+            "test.zips.exr", window, window, dataZ, dataA, sampleCount);
+    }
 }
diff --git a/src/examples/deepTiledExamples.cpp b/src/examples/deepTiledExamples.cpp
@@ -163,7 +163,8 @@ writeDeepTiledFile (
     Box2i      displayWindow,
     Box2i      dataWindow,
     int        tileSizeX,
-    int        tileSizeY)
+    int        tileSizeY,
+    Compression compression = Compression::ZIPS_COMPRESSION)
 {
     //
     // Write a deep image with only a A (alpha) and a Z (depth) channel,
@@ -183,7 +184,7 @@ writeDeepTiledFile (
     header.channels ().insert ("Z", Channel (FLOAT));
     header.channels ().insert ("A", Channel (HALF));
     header.setType (DEEPTILE);
-    header.compression () = ZIPS_COMPRESSION;
+    header.compression () = compression;
 
     header.setTileDescription (
         TileDescription (tileSizeX, tileSizeY, ONE_LEVEL));
@@ -273,6 +274,8 @@ void deepTiledExamples()
     testDataZ.resizeErase(h, w);
     drawImage2(testDataA, testDataZ, w, h);
 
-    writeDeepTiledFile("testTiled.deep.exr", window, window, tileSizeX, tileSizeY);
-    readDeepTiledFile ("testTiled.deep.exr", window, window, dataZ, dataA, sampleCount);
+    writeDeepTiledFile("testTiled.deep.zip.exr", window, window, tileSizeX, tileSizeY);
+    readDeepTiledFile ("testTiled.deep.zip.exr", window, window, dataZ, dataA, sampleCount);
+    writeDeepTiledFile("testTiled.deep.zstd.exr", window, window, tileSizeX, tileSizeY, Compression::ZSTD_COMPRESSION);
+    readDeepTiledFile ("testTiled.deep.zstd.exr", window, window, dataZ, dataA, sampleCount);
 }
diff --git a/src/lib/OpenEXR/CMakeLists.txt b/src/lib/OpenEXR/CMakeLists.txt
@@ -29,6 +29,7 @@ openexr_define_library(OpenEXR
     ImfTiledMisc.h
     ImfZip.h
     ImfZipCompressor.h
+    ImfZstdCompressor.h
     b44ExpLogTable.h
     dwaLookups.h
     ImfAcesFile.cpp
@@ -122,6 +123,7 @@ openexr_define_library(OpenEXR
     ImfWav.cpp
     ImfZip.cpp
     ImfZipCompressor.cpp
+    ImfZstdCompressor.cpp
   HEADERS
     ImfAcesFile.h
     ImfArray.h
@@ -220,3 +222,7 @@ openexr_define_library(OpenEXR
     OpenEXR::IlmThread
     OpenEXR::OpenEXRCore
   )
+
+target_include_directories(OpenEXR PUBLIC "/home/vladal/bin/include/")
+target_link_directories(OpenEXR PUBLIC "/home/vladal/bin/lib")
+target_link_libraries(OpenEXR PUBLIC "dl" "blosc2")
diff --git a/src/lib/OpenEXR/ImfCompression.h b/src/lib/OpenEXR/ImfCompression.h
@@ -45,6 +45,8 @@ enum IMF_EXPORT_ENUM Compression
                           // wise and faster to decode full frames
                           // than DWAA_COMPRESSION.
 
+    ZSTD_COMPRESSION = 10,
+
     NUM_COMPRESSION_METHODS // number of different compression methods
 };
 

diff --git a/src/lib/OpenEXR/ImfCompressionAttribute.cpp b/src/lib/OpenEXR/ImfCompressionAttribute.cpp
@@ -57,7 +57,8 @@ CompressionAttribute::readValueFrom (
         tmp != ZIPS_COMPRESSION && tmp != ZIP_COMPRESSION &&
         tmp != PIZ_COMPRESSION && tmp != PXR24_COMPRESSION &&
         tmp != B44_COMPRESSION && tmp != B44A_COMPRESSION &&
-        tmp != DWAA_COMPRESSION && tmp != DWAB_COMPRESSION)
+        tmp != DWAA_COMPRESSION && tmp != DWAB_COMPRESSION && 
+        tmp != ZSTD_COMPRESSION)
     {
         tmp = NUM_COMPRESSION_METHODS;
     }

diff --git a/src/lib/OpenEXR/ImfCompressor.cpp b/src/lib/OpenEXR/ImfCompressor.cpp
@@ -18,6 +18,7 @@
 #include "ImfPxr24Compressor.h"
 #include "ImfRleCompressor.h"
 #include "ImfZipCompressor.h"
+#include "ImfZstdCompressor.h"
 
 OPENEXR_IMF_INTERNAL_NAMESPACE_SOURCE_ENTER
 
@@ -63,7 +64,8 @@ isValidCompression (Compression c)
         case B44_COMPRESSION:
         case B44A_COMPRESSION:
         case DWAA_COMPRESSION:
-        case DWAB_COMPRESSION: return true;
+        case DWAB_COMPRESSION:
+        case ZSTD_COMPRESSION: return true;
 
         default: return false;
     }
@@ -89,7 +91,8 @@ isValidDeepCompression (Compression c)
     {
         case NO_COMPRESSION:
         case RLE_COMPRESSION:
-        case ZIPS_COMPRESSION: return true;
+        case ZIPS_COMPRESSION:
+        case ZSTD_COMPRESSION: return true;
         default: return false;
     }
 }
@@ -141,6 +144,8 @@ newCompressor (Compression c, size_t maxScanLineSize, const Header& hdr)
                 256,
                 DwaCompressor::STATIC_HUFFMAN);
 
+        case ZSTD_COMPRESSION:
+            return new ZstdCompressor (hdr, maxScanLineSize, 32);
         default: return 0;
     }
 }
@@ -162,6 +167,7 @@ numLinesInBuffer (Compression comp)
         case B44_COMPRESSION:
         case B44A_COMPRESSION:
         case DWAA_COMPRESSION: return 32;
+        case ZSTD_COMPRESSION: return 32;
         case DWAB_COMPRESSION: return 256;
 
         default: throw IEX_NAMESPACE::ArgExc ("Unknown compression type");
@@ -182,6 +188,9 @@ newTileCompressor (
         case ZIP_COMPRESSION:
 
             return new ZipCompressor (hdr, tileLineSize, numTileLines);
+        case ZSTD_COMPRESSION:
+
+            return new ZstdCompressor (hdr, tileLineSize, numTileLines);
 
         case PIZ_COMPRESSION:
 

diff --git a/src/lib/OpenEXR/ImfHeader.cpp b/src/lib/OpenEXR/ImfHeader.cpp
@@ -70,9 +70,11 @@ struct CompressionRecord
     {
         exr_get_default_zip_compression_level(&zip_level);
         exr_get_default_dwa_compression_quality(&dwa_level);
+        exr_get_default_zstd_compression_level(&zstd_level);
     }
     int   zip_level;
     float dwa_level;
+    int   zstd_level;
 };
 // NB: This is extra complicated than one would normally write to
 // handle scenario that seems to happen on MacOS/Windows (probably
@@ -696,6 +698,18 @@ Header::zipCompressionLevel () const
     return retrieveCompressionRecord (this).zip_level;
 }
 
+int&
+Header::zstdCompressionLevel ()
+{
+    return retrieveCompressionRecord (this).zstd_level;
+}
+
+int
+Header::zstdCompressionLevel () const
+{
+    return retrieveCompressionRecord (this).zstd_level;
+}
+
 float&
 Header::dwaCompressionLevel ()
 {

diff --git a/src/lib/OpenEXR/ImfHeader.h b/src/lib/OpenEXR/ImfHeader.h
@@ -285,6 +285,11 @@ class IMF_EXPORT_TYPE Header
     float& dwaCompressionLevel ();
     IMF_EXPORT
     float dwaCompressionLevel () const;
+    IMF_EXPORT
+    int& zstdCompressionLevel ();
+    IMF_EXPORT
+    int zstdCompressionLevel () const;
+    IMF_EXPORT
 
     //-----------------------------------------------------
     // Access to required attributes for multipart files

diff --git a/src/lib/OpenEXR/ImfMultiPartInputFile.cpp b/src/lib/OpenEXR/ImfMultiPartInputFile.cpp
@@ -547,6 +547,7 @@ MultiPartInputFile::Data::chunkOffsetReconstruction (
             // (TODO) fix this so that it doesn't need to be revised for future compression types.
             switch (parts[i]->header.compression ())
             {
+                case ZSTD_COMPRESSION: rowsizes[i] = 32; break;
                 case DWAB_COMPRESSION: rowsizes[i] = 256; break;
                 case PIZ_COMPRESSION:
                 case B44_COMPRESSION:
-Original file line number
+Diff line change
@@ Expand Up / @@ -57,3 +57,7 @@ docs/_test_images/ @@
     # Ignore Bazel generated files
     bazel-*
+    cmake-build-debug/
+    .idea/