From b936ded35e13e6f68d27a4b0768fb158a655f369 Mon Sep 17 00:00:00 2001
From: Vojtech Moravec <vojtech.moravec.st@vsb.cz>
Date: Fri, 24 Jan 2020 09:10:12 +0100
Subject: [PATCH] Support codebook from cache files.

---
 .../java/azgracompress/DataCompressor.java    |  7 ++++
 .../java/azgracompress/cli/CliConstants.java  |  3 ++
 .../azgracompress/cli/ParsedCliOptions.java   | 15 ++++++++-
 .../compression/SQImageCompressor.java        | 29 ++++++++++++++--
 .../compression/VQImageCompressor.java        | 33 +++++++++++++++++--
 .../quantization/QuantizationValueCache.java  | 17 ++++------
 6 files changed, 86 insertions(+), 18 deletions(-)

diff --git a/src/main/java/azgracompress/DataCompressor.java b/src/main/java/azgracompress/DataCompressor.java
index ec68519..22d5aab 100644
--- a/src/main/java/azgracompress/DataCompressor.java
+++ b/src/main/java/azgracompress/DataCompressor.java
@@ -131,10 +131,12 @@ public class DataCompressor {
         options.addOptionGroup(methodGroup);
         options.addOptionGroup(compressionMethodGroup);
         options.addOption(CliConstants.BITS_SHORT, CliConstants.BITS_LONG, true, "Bit count per pixel [Default 8]");
+
         options.addOption(CliConstants.REFERENCE_PLANE_SHORT,
                           CliConstants.REFERENCE_PLANE_LONG,
                           true,
                           "Reference plane index");
+
         options.addOption(new Option(CliConstants.VERBOSE_SHORT,
                                      CliConstants.VERBOSE_LONG,
                                      false,
@@ -145,6 +147,11 @@ public class DataCompressor {
                                      true,
                                      "Number of worker threads"));
 
+        options.addOption(new Option(CliConstants.CODEBOOK_CACHE_FOLDER_SHORT,
+                                     CliConstants.CODEBOOK_CACHE_FOLDER_LONG,
+                                     true,
+                                     "Folder of codebook caches"));
+
         options.addOption(CliConstants.OUTPUT_SHORT, CliConstants.OUTPUT_LONG, true, "Custom output file");
         return options;
     }
diff --git a/src/main/java/azgracompress/cli/CliConstants.java b/src/main/java/azgracompress/cli/CliConstants.java
index 644e300..29222fd 100644
--- a/src/main/java/azgracompress/cli/CliConstants.java
+++ b/src/main/java/azgracompress/cli/CliConstants.java
@@ -34,6 +34,9 @@ public class CliConstants {
     public static final String WORKER_COUNT_SHORT = "wc";
     public static final String WORKER_COUNT_LONG = "worker-count";
 
+    public static final String CODEBOOK_CACHE_FOLDER_SHORT = "wc";
+    public static final String CODEBOOK_CACHE_FOLDER_LONG = "worker-count";
+
     public static final String SCALAR_QUANTIZATION_SHORT = "sq";
     public static final String SCALAR_QUANTIZATION_LONG = "scalar-quantization";
 
diff --git a/src/main/java/azgracompress/cli/ParsedCliOptions.java b/src/main/java/azgracompress/cli/ParsedCliOptions.java
index d60f488..c208e5c 100644
--- a/src/main/java/azgracompress/cli/ParsedCliOptions.java
+++ b/src/main/java/azgracompress/cli/ParsedCliOptions.java
@@ -18,6 +18,7 @@ public class ParsedCliOptions {
 
     private String inputFile;
     private String outputFile;
+    private String codebookCacheFolder = null;
 
     private int bitsPerPixel;
     private V2i vectorDimension = new V2i(0);
@@ -107,6 +108,8 @@ public class ParsedCliOptions {
             }
         }
 
+        codebookCacheFolder = cmd.getOptionValue(CliConstants.CODEBOOK_CACHE_FOLDER_LONG, null);
+
         if (!errorOccurred) {
             outputFile = cmd.getOptionValue(CliConstants.OUTPUT_LONG, getDefaultOutputFilePath(inputFile));
         }
@@ -404,6 +407,14 @@ public class ParsedCliOptions {
         return workerCount;
     }
 
+    public String getCodebookCacheFolder() {
+        return codebookCacheFolder;
+    }
+
+    public boolean hasCodebookCacheFolder() {
+        return (codebookCacheFolder != null);
+    }
+
     public String report() {
         StringBuilder sb = new StringBuilder();
 
@@ -443,6 +454,9 @@ public class ParsedCliOptions {
         sb.append("BitsPerPixel: ").append(bitsPerPixel).append('\n');
         sb.append("Output: ").append(outputFile).append('\n');
         sb.append("InputFile: ").append(inputFile).append('\n');
+        if (hasCodebookCacheFolder()) {
+            sb.append("CodebookCacheFolder: ").append(codebookCacheFolder).append('\n');
+        }
 
         if (hasQuantizationType(method)) {
             sb.append("Input image dims: ").append(imageDimension.toString()).append('\n');
@@ -462,7 +476,6 @@ public class ParsedCliOptions {
         sb.append("Verbose: ").append(verbose).append('\n');
         sb.append("ThreadWorkerCount: ").append(workerCount).append('\n');
 
-
         return sb.toString();
     }
 
diff --git a/src/main/java/azgracompress/compression/SQImageCompressor.java b/src/main/java/azgracompress/compression/SQImageCompressor.java
index d0461ab..f49727d 100644
--- a/src/main/java/azgracompress/compression/SQImageCompressor.java
+++ b/src/main/java/azgracompress/compression/SQImageCompressor.java
@@ -54,6 +54,22 @@ public class SQImageCompressor extends CompressorDecompressorBase implements IIm
         }
     }
 
+    /**
+     * Load quantization codebook from cache file.
+     *
+     * @return Scalar quantizer with cached codebook.
+     * @throws ImageCompressionException when fails to read cached codebook.
+     */
+    private ScalarQuantizer loadQuantizerFromCache() throws ImageCompressionException {
+        QuantizationValueCache cache = new QuantizationValueCache(options.getCodebookCacheFolder());
+        try {
+            final int[] quantizationValues = cache.readCachedValues(options.getInputFile(), codebookSize);
+            return new ScalarQuantizer(U16.Min, U16.Max, quantizationValues);
+        } catch (IOException e) {
+            throw new ImageCompressionException("Failed to read quantization values from cache file.", e);
+        }
+    }
+
     /**
      * Compress the image file specified by parsed CLI options using scalar quantization.
      *
@@ -61,9 +77,16 @@ public class SQImageCompressor extends CompressorDecompressorBase implements IIm
      * @throws ImageCompressionException When compress process fails.
      */
     public void compress(DataOutputStream compressStream) throws ImageCompressionException {
-        ScalarQuantizer quantizer = null;
         Stopwatch stopwatch = new Stopwatch();
-        if (options.hasReferencePlaneIndex()) {
+        final boolean hasGeneralQuantizer = options.hasCodebookCacheFolder() || options.hasReferencePlaneIndex();
+
+        ScalarQuantizer quantizer = null;
+        if (options.hasCodebookCacheFolder()) {
+            Log("Loading codebook from cache file.");
+            quantizer = loadQuantizerFromCache();
+            Log("Cached quantizer created.");
+        } else if (options.hasReferencePlaneIndex()) {
+            // TODO(Moravec): Reference plane will be deprecated in favor of 'middle' plane.
             stopwatch.restart();
             ImageU16 referencePlane = null;
             try {
@@ -99,7 +122,7 @@ public class SQImageCompressor extends CompressorDecompressorBase implements IIm
                 throw new ImageCompressionException("Unable to load plane data.", ex);
             }
 
-            if (!options.hasReferencePlaneIndex()) {
+            if (!hasGeneralQuantizer) {
                 Log(String.format("Training scalar quantizer from plane %d.", planeIndex));
                 quantizer = trainScalarQuantizerFromData(plane.getData());
                 writeCodebookToOutputStream(quantizer, compressStream);
diff --git a/src/main/java/azgracompress/compression/VQImageCompressor.java b/src/main/java/azgracompress/compression/VQImageCompressor.java
index 3ad1972..e2ea3ce 100644
--- a/src/main/java/azgracompress/compression/VQImageCompressor.java
+++ b/src/main/java/azgracompress/compression/VQImageCompressor.java
@@ -69,6 +69,26 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
         }
     }
 
+    /**
+     * Load quantizer from cached codebook.
+     *
+     * @return Vector quantizer with cached codebook.
+     * @throws ImageCompressionException when fails to read cached codebook.
+     */
+    private VectorQuantizer loadQuantizerFromCache() throws ImageCompressionException {
+        QuantizationValueCache cache = new QuantizationValueCache(options.getCodebookCacheFolder());
+        try {
+            final CodebookEntry[] codebook = cache.readCachedValues(options.getInputFile(),
+                                                                    codebookSize,
+                                                                    options.getVectorDimension().getX(),
+                                                                    options.getVectorDimension().getY());
+            return new VectorQuantizer(codebook);
+
+        } catch (IOException e) {
+            throw new ImageCompressionException("Failed to read quantization vectors from cache.", e);
+        }
+    }
+
     /**
      * Compress the image file specified by parsed CLI options using vector quantization.
      *
@@ -76,9 +96,15 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
      * @throws ImageCompressionException When compress process fails.
      */
     public void compress(DataOutputStream compressStream) throws ImageCompressionException {
-        VectorQuantizer quantizer = null;
         Stopwatch stopwatch = new Stopwatch();
-        if (options.hasReferencePlaneIndex()) {
+        final boolean hasGeneralQuantizer = options.hasCodebookCacheFolder() || options.hasReferencePlaneIndex();
+        VectorQuantizer quantizer = null;
+
+        if (options.hasCodebookCacheFolder()) {
+            Log("Loading codebook from cache file.");
+            quantizer = loadQuantizerFromCache();
+            Log("Cached quantizer created.");
+        } else if (options.hasReferencePlaneIndex()) {
             stopwatch.restart();
 
             ImageU16 referencePlane = null;
@@ -116,7 +142,7 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
 
             final int[][] planeVectors = getPlaneVectors(plane);
 
-            if (!options.hasReferencePlaneIndex()) {
+            if (!hasGeneralQuantizer) {
                 Log(String.format("Training vector quantizer from plane %d.", planeIndex));
                 quantizer = trainVectorQuantizerFromPlaneVectors(planeVectors);
                 writeQuantizerToCompressStream(quantizer, compressStream);
@@ -139,6 +165,7 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
         }
     }
 
+
     /**
      * Load plane and convert the plane into quantization vectors.
      *
diff --git a/src/main/java/azgracompress/quantization/QuantizationValueCache.java b/src/main/java/azgracompress/quantization/QuantizationValueCache.java
index 87a4972..fca4c95 100644
--- a/src/main/java/azgracompress/quantization/QuantizationValueCache.java
+++ b/src/main/java/azgracompress/quantization/QuantizationValueCache.java
@@ -26,7 +26,10 @@ public class QuantizationValueCache {
                                              final int entryHeight) {
         final File inputFile = new File(trainFile);
         final File cacheFile = new File(cacheFolder, String.format("%s_%d_%dx%d.qvc",
-                                                                   inputFile.getName(), codebookSize, entryWidth, entryHeight));
+                                                                   inputFile.getName(),
+                                                                   codebookSize,
+                                                                   entryWidth,
+                                                                   entryHeight));
         return cacheFile;
     }
 
@@ -72,7 +75,7 @@ public class QuantizationValueCache {
         }
     }
 
-    public int[] readCachedValues(final String trainFile, final int quantizationValueCount) {
+    public int[] readCachedValues(final String trainFile, final int quantizationValueCount) throws IOException {
         final File cacheFile = getCacheFileForScalarValues(trainFile, quantizationValueCount);
 
         int[] values = new int[quantizationValueCount];
@@ -82,10 +85,6 @@ public class QuantizationValueCache {
             for (int i = 0; i < quantizationValueCount; i++) {
                 values[i] = dis.readInt();
             }
-        } catch (IOException ioEx) {
-            System.err.println("Failed to read scalar quantization values from cache.");
-            ioEx.printStackTrace();
-            return new int[0];
         }
         return values;
     }
@@ -93,7 +92,7 @@ public class QuantizationValueCache {
     public CodebookEntry[] readCachedValues(final String trainFile,
                                             final int codebookSize,
                                             final int entryWidth,
-                                            final int entryHeight) {
+                                            final int entryHeight) throws IOException {
         final File cacheFile = getCacheFileForVectorValues(trainFile, codebookSize, entryWidth, entryHeight);
 
         CodebookEntry[] codebook = new CodebookEntry[codebookSize];
@@ -115,10 +114,6 @@ public class QuantizationValueCache {
                 }
                 codebook[i] = new CodebookEntry(vector);
             }
-        } catch (IOException ioEx) {
-            System.err.println("Failed to read quantization vectors from cache.");
-            ioEx.printStackTrace();
-            return new CodebookEntry[0];
         }
         return codebook;
 
-- 
GitLab