From 254b3048f45fd70510cfeb7e82f6e35a26bc3f56 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20Moravec?= <theazgra@gmail.com>
Date: Thu, 23 Jul 2020 12:33:53 +0200
Subject: [PATCH] Extend caching mechanism to support plugin.

Cache managers need the cache folder to be set.
In the CLI app we set the cache folder to the output file as it was
previously.

Caching functions also requires file name of caching hing, which is used
when saving the cache file itself. BufferInputData didn't have this
hing, so we had to add it. Cache file hint is requested by the function
call.
---
 .../cache/QuantizationCacheManager.java       | 40 +++++++++++++++----
 .../azgracompress/cli/ParsedCliOptions.java   |  6 ++-
 .../compression/SQImageCompressor.java        | 11 +++--
 .../compression/VQImageCompressor.java        | 24 ++++++-----
 .../azgracompress/io/BufferInputData.java     | 13 +++++-
 .../java/azgracompress/io/FileInputData.java  |  7 ++++
 src/main/java/azgracompress/io/InputData.java |  8 ++++
 7 files changed, 87 insertions(+), 22 deletions(-)

diff --git a/src/main/java/azgracompress/cache/QuantizationCacheManager.java b/src/main/java/azgracompress/cache/QuantizationCacheManager.java
index dbd07c8..d17d3ef 100644
--- a/src/main/java/azgracompress/cache/QuantizationCacheManager.java
+++ b/src/main/java/azgracompress/cache/QuantizationCacheManager.java
@@ -2,6 +2,7 @@ package azgracompress.cache;
 
 import azgracompress.data.V3i;
 import azgracompress.fileformat.QuantizationType;
+import azgracompress.quantization.QTrainIteration;
 import azgracompress.quantization.scalar.SQCodebook;
 import azgracompress.quantization.vector.VQCodebook;
 
@@ -36,7 +37,7 @@ public class QuantizationCacheManager {
     private File getCacheFilePathForSQ(final String trainFile, final int codebookSize) {
         final File inputFile = new File(trainFile);
         return new File(cacheFolder, String.format("%s_%d_bits.qvc",
-                                                   inputFile.getName(), codebookSize));
+                inputFile.getName(), codebookSize));
     }
 
     /**
@@ -52,7 +53,7 @@ public class QuantizationCacheManager {
                                        final V3i vDim) {
         final File inputFile = new File(trainFile);
         return new File(cacheFolder, String.format("%s_%d_%dx%d.qvc", inputFile.getName(), codebookSize,
-                                                   vDim.getX(), vDim.getY()));
+                vDim.getX(), vDim.getY()));
     }
 
 
@@ -140,8 +141,8 @@ public class QuantizationCacheManager {
      */
     public void saveCodebook(final String trainFile, final VQCodebook codebook) throws IOException {
         final String fileName = getCacheFilePathForVQ(trainFile,
-                                                      codebook.getCodebookSize(),
-                                                      codebook.getVectorDims()).getAbsolutePath();
+                codebook.getCodebookSize(),
+                codebook.getVectorDims()).getAbsolutePath();
 
         final CacheFileHeader header = createHeaderForVQ(new File(trainFile).getName(), codebook);
         final VQCacheFile cacheFile = new VQCacheFile(header, codebook);
@@ -172,15 +173,38 @@ public class QuantizationCacheManager {
         }
     }
 
+    /**
+     * Check if the SQ cache file for given image file exists.
+     *
+     * @param imageFile    Image file.
+     * @param codebookSize Scalar quantization codebook size.
+     * @return True if cache file exists and and can be loaded.
+     */
+    public boolean doesSQCacheExists(final String imageFile, final int codebookSize) {
+        return getCacheFilePathForSQ(imageFile, codebookSize).exists();
+    }
+
+    /**
+     * Check if the VQ cache file for given image file exists.
+     *
+     * @param imageFile    Image file.
+     * @param codebookSize Scalar quantization codebook size.
+     * @param vDim         Quantization vector dimensions.
+     * @return True if cache file exists and and can be loaded.
+     */
+    public boolean doesVQCacheExists(final String imageFile, final int codebookSize, final V3i vDim) {
+        return getCacheFilePathForVQ(imageFile, codebookSize, vDim).exists();
+    }
+
     /**
      * Load SQ cache file from disk.
      *
-     * @param trainFile    Input image file.
+     * @param imageFile    Input image file.
      * @param codebookSize Codebook size.
      * @return SQ cache file.
      */
-    private SQCacheFile loadSQCacheFile(final String trainFile, final int codebookSize) {
-        final File path = getCacheFilePathForSQ(trainFile, codebookSize);
+    private SQCacheFile loadSQCacheFile(final String imageFile, final int codebookSize) {
+        final File path = getCacheFilePathForSQ(imageFile, codebookSize);
         try {
             return (SQCacheFile) readCacheFile(path, new SQCacheFile());
         } catch (IOException e) {
@@ -293,4 +317,6 @@ public class QuantizationCacheManager {
         sb.append('\n');
         System.out.println(sb.toString());
     }
+
+
 }
diff --git a/src/main/java/azgracompress/cli/ParsedCliOptions.java b/src/main/java/azgracompress/cli/ParsedCliOptions.java
index 1ff3dfe..d46ead0 100644
--- a/src/main/java/azgracompress/cli/ParsedCliOptions.java
+++ b/src/main/java/azgracompress/cli/ParsedCliOptions.java
@@ -149,10 +149,14 @@ public class ParsedCliOptions extends CompressionOptions implements Cloneable {
         }
 
         if (!parseErrorOccurred) {
-            setOutputFilePath(cmd.getOptionValue(CliConstants.OUTPUT_LONG, getDefaultOutputFilePath(((FileInputData) getInputDataInfo()).getFilePath())));
+            setOutputFilePath(cmd.getOptionValue(CliConstants.OUTPUT_LONG, getDefaultOutputFilePath(getInputDataInfo().getFilePath())));
             setCodebookCacheFolder(cmd.getOptionValue(CliConstants.CODEBOOK_CACHE_FOLDER_LONG, null));
         }
 
+        if (getMethod() == ProgramMethod.TrainCodebook) {
+            setCodebookCacheFolder(getOutputFilePath());
+        }
+
         parseError = errorBuilder.toString();
     }
 
diff --git a/src/main/java/azgracompress/compression/SQImageCompressor.java b/src/main/java/azgracompress/compression/SQImageCompressor.java
index 598f231..7a3b067 100644
--- a/src/main/java/azgracompress/compression/SQImageCompressor.java
+++ b/src/main/java/azgracompress/compression/SQImageCompressor.java
@@ -73,8 +73,11 @@ public class SQImageCompressor extends CompressorDecompressorBase implements IIm
     private ScalarQuantizer loadQuantizerFromCache() throws ImageCompressionException {
         QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getCodebookCacheFolder());
 
-        final SQCodebook codebook = cacheManager.loadSQCodebook(options.getInputDataInfo().getFilePath(),
-                getCodebookSize());
+        if (!cacheManager.doesSQCacheExists(options.getInputDataInfo().getCacheFileName(), getCodebookSize())) {
+            trainAndSaveCodebook();
+        }
+
+        final SQCodebook codebook = cacheManager.loadSQCodebook(options.getInputDataInfo().getCacheFileName(), getCodebookSize());
         if (codebook == null) {
             throw new ImageCompressionException("Failed to read quantization values from cache file.");
         }
@@ -220,9 +223,9 @@ public class SQImageCompressor extends CompressorDecompressorBase implements IIm
         Log("Finished LloydMax training.");
 
         Log(String.format("Saving cache file to %s", options.getOutputFilePath()));
-        QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getOutputFilePath());
+        QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getCodebookCacheFolder());
         try {
-            cacheManager.saveCodebook(options.getInputDataInfo().getFilePath(), codebook);
+            cacheManager.saveCodebook(options.getInputDataInfo().getCacheFileName(), codebook);
         } catch (IOException e) {
             throw new ImageCompressionException("Unable to write cache.", e);
         }
diff --git a/src/main/java/azgracompress/compression/VQImageCompressor.java b/src/main/java/azgracompress/compression/VQImageCompressor.java
index e1f89b4..d94984c 100644
--- a/src/main/java/azgracompress/compression/VQImageCompressor.java
+++ b/src/main/java/azgracompress/compression/VQImageCompressor.java
@@ -75,10 +75,16 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
 
         QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getCodebookCacheFolder());
 
+        if (!cacheManager.doesVQCacheExists(options.getInputDataInfo().getCacheFileName(),
+                getCodebookSize(),
+                options.getVectorDimension().toV3i())) {
+            trainAndSaveCodebook();
+        }
+
+        final VQCodebook codebook = cacheManager.loadVQCodebook(options.getInputDataInfo().getCacheFileName(),
+                getCodebookSize(),
+                options.getVectorDimension().toV3i());
 
-        final VQCodebook codebook = cacheManager.loadVQCodebook(options.getInputDataInfo().getFilePath(),
-                                                                getCodebookSize(),
-                                                                options.getVectorDimension().toV3i());
         if (codebook == null) {
             throw new ImageCompressionException("Failed to read quantization vectors from cache.");
         }
@@ -230,10 +236,10 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
                 }
 
                 System.arraycopy(planeVectors,
-                                 0,
-                                 trainData,
-                                 (planeCounter * chunkCountPerPlane),
-                                 chunkCountPerPlane);
+                        0,
+                        trainData,
+                        (planeCounter * chunkCountPerPlane),
+                        chunkCountPerPlane);
                 ++planeCounter;
             }
         }
@@ -256,9 +262,9 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
 
 
         Log("Saving cache file to %s", options.getOutputFilePath());
-        QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getOutputFilePath());
+        QuantizationCacheManager cacheManager = new QuantizationCacheManager(options.getCodebookCacheFolder());
         try {
-            cacheManager.saveCodebook(options.getInputDataInfo().getFilePath(), lbgResult.getCodebook());
+            cacheManager.saveCodebook(options.getInputDataInfo().getCacheFileName(), lbgResult.getCodebook());
         } catch (IOException e) {
             throw new ImageCompressionException("Unable to write VQ cache.", e);
         }
diff --git a/src/main/java/azgracompress/io/BufferInputData.java b/src/main/java/azgracompress/io/BufferInputData.java
index 9ab179d..3d0223a 100644
--- a/src/main/java/azgracompress/io/BufferInputData.java
+++ b/src/main/java/azgracompress/io/BufferInputData.java
@@ -12,6 +12,7 @@ public class BufferInputData extends InputData {
      */
     private final Object[] imageBuffers;
 
+    private final String cacheHint;
 
     /**
      * Create input data backed by buffer object.
@@ -19,12 +20,17 @@ public class BufferInputData extends InputData {
      * @param imageBuffers    Image buffer references.
      * @param imageDimensions Image dimensions.
      * @param pixelType       Image pixel type.
+     * @param cacheHint       Name of the image used in caching.
      */
-    public BufferInputData(final Object[] imageBuffers, final V3i imageDimensions, final PixelType pixelType) {
+    public BufferInputData(final Object[] imageBuffers,
+                           final V3i imageDimensions,
+                           final PixelType pixelType,
+                           final String cacheHint) {
         this.imageBuffers = imageBuffers;
         setDataLoaderType(DataLoaderType.ImageJBufferLoader);
         setDimension(imageDimensions);
         setPixelType(pixelType);
+        this.cacheHint = cacheHint;
     }
 
     /**
@@ -38,4 +44,9 @@ public class BufferInputData extends InputData {
         assert (planeIndex < imageBuffers.length);
         return imageBuffers[planeIndex];
     }
+
+    @Override
+    public String getCacheFileName() {
+        return cacheHint;
+    }
 }
diff --git a/src/main/java/azgracompress/io/FileInputData.java b/src/main/java/azgracompress/io/FileInputData.java
index 49fad66..968fbd6 100644
--- a/src/main/java/azgracompress/io/FileInputData.java
+++ b/src/main/java/azgracompress/io/FileInputData.java
@@ -12,6 +12,7 @@ public class FileInputData extends InputData {
 
     /**
      * Create input data backed by data file.
+     *
      * @param filePath
      */
     public FileInputData(String filePath) {
@@ -20,10 +21,16 @@ public class FileInputData extends InputData {
 
     /**
      * Get path to the data file.
+     *
      * @return
      */
     @Override
     public String getFilePath() {
         return filePath;
     }
+
+    @Override
+    public String getCacheFileName() {
+        return filePath;
+    }
 }
diff --git a/src/main/java/azgracompress/io/InputData.java b/src/main/java/azgracompress/io/InputData.java
index f4a27a9..966a706 100644
--- a/src/main/java/azgracompress/io/InputData.java
+++ b/src/main/java/azgracompress/io/InputData.java
@@ -97,9 +97,17 @@ public abstract class InputData {
 
     /**
      * Override in FileInputData!!!
+     *
      * @return null!
      */
     public String getFilePath() {
         return null;
     }
+
+    /**
+     * Get name used in creation of qcmp cache file.
+     *
+     * @return Name used for cache file.
+     */
+    public abstract String getCacheFileName();
 }
-- 
GitLab