diff --git a/src/main/java/cz/it4i/qcmp/compression/ImageCompressor.java b/src/main/java/cz/it4i/qcmp/compression/ImageCompressor.java
index ac87edb777610c2db1bdbb03d9c10ea27db46193..61d5d45b5fdcd5e83f04004589f818736156d911 100644
--- a/src/main/java/cz/it4i/qcmp/compression/ImageCompressor.java
+++ b/src/main/java/cz/it4i/qcmp/compression/ImageCompressor.java
@@ -197,6 +197,19 @@ public class ImageCompressor extends CompressorDecompressorBase {
         }
     }
 
+    /**
+     * Allow usage of the KD tree when looking for the best quantization vector.
+     *
+     * @return True if imageCompressor is VQImageCompressor and KD tree was enabled.
+     */
+    public boolean allowKdTreeVectorLookup() {
+        if ((imageCompressor != null) && (imageCompressor instanceof VQImageCompressor)) {
+            ((VQImageCompressor) imageCompressor).setUseKdTree(true);
+            return true;
+        }
+        return false;
+    }
+
 
     /**
      * Create QCMPFile header for compressed file.
diff --git a/src/main/java/cz/it4i/qcmp/compression/VQImageCompressor.java b/src/main/java/cz/it4i/qcmp/compression/VQImageCompressor.java
index ae130701b39f54d7fa3123dc9e4c34c644842c58..05bb5fe80b2a0a39ef9c9965ad0b11bff89d4f88 100644
--- a/src/main/java/cz/it4i/qcmp/compression/VQImageCompressor.java
+++ b/src/main/java/cz/it4i/qcmp/compression/VQImageCompressor.java
@@ -26,6 +26,8 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
     private VectorQuantizer cachedQuantizer = null;
     private Huffman cachedHuffman = null;
 
+    private boolean useKdTree = false;
+
     public VQImageCompressor(final CompressionOptions options) {
         super(options);
     }
@@ -37,6 +39,14 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
         cachedHuffman = createHuffmanCoder(createHuffmanSymbols(cachedCodebook.getCodebookSize()), cachedCodebook.getVectorFrequencies());
     }
 
+    public boolean shouldUseKdTree() {
+        return useKdTree;
+    }
+
+    public void setUseKdTree(final boolean useKdTree) {
+        this.useKdTree = useKdTree;
+    }
+
     /**
      * Train vector quantizer from plane vectors.
      *
@@ -200,10 +210,7 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
                 writeQuantizerToCompressStream(quantizer, compressStream);
             }
 
-            // Use BestBinFirst KDTree for codebook lookup.
-            // final int[] indices = quantizer.quantizeIntoIndicesUsingKDTree(planeVectors, options.getWorkerCount());
-            // Use BruteForce for codebook lookup.
-            final int[] indices = quantizer.quantizeIntoIndices(planeVectors, options.getWorkerCount());
+            final int[] indices = quantizeVectorsImpl(quantizer, planeVectors, options.getWorkerCount());
 
             planeDataSizes[planeCounter++] = writeHuffmanEncodedIndices(compressStream, huffman, indices);
 
@@ -219,6 +226,20 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
         return planeDataSizes;
     }
 
+    /**
+     * Quantize into indices call wrapper for KDTree switch.
+     *
+     * @param quantizer   Vector quantizer.
+     * @param srcVectors  Vectors to quantize.
+     * @param workerCount Maximum number of worker threads.
+     * @return Indices of codebook vectors.
+     */
+    private int[] quantizeVectorsImpl(final VectorQuantizer quantizer, final int[][] srcVectors, final int workerCount) {
+        if (useKdTree)
+            return quantizer.quantizeIntoIndicesUsingKDTree(srcVectors, workerCount);
+        return quantizer.quantizeIntoIndices(srcVectors, workerCount);
+    }
+
     /**
      * Calculate the number of voxel layers needed for dataset of plane count.
      *
@@ -284,8 +305,8 @@ public class VQImageCompressor extends CompressorDecompressorBase implements IIm
                 throw new ImageCompressionException("Unable to load voxels from voxel layer " + voxelLayerRange, e);
             }
 
-            final int[] indices = quantizer.quantizeIntoIndices(voxelData, options.getWorkerCount());
-            // final int[] indices = quantizer.quantizeIntoIndicesUsingKDTree(voxelData, options.getWorkerCount());
+            final int[] indices = quantizeVectorsImpl(quantizer, voxelData, options.getWorkerCount());
+            
             voxelLayersSizes[voxelLayerIndex] = writeHuffmanEncodedIndices(compressStream, huffman, indices);
             stopwatch.stop();
             if (options.isConsoleApplication()) {