feat: add pipeline upper bound

Before the upper bound was an argument, now it is defined with the pipeline, and it generalizes to any pipeline

antecedent/block.py

@@ -35,7 +35,6 @@ class Block:
     def search_antecedent(self,
                           pipeline: ComposedPipeline,
                           max_iter: int,
-                          upper_bound: np.ndarray = None,
                           shared_dict: dict = None,
                           task_id: int = None,
                           verbose: bool = False):
@@ -46,7 +45,6 @@ class Block:
             pipeline: must have a forward_dct, an inverse_dct, a round_pixel and a round_dct method
             max_iter: maximum number of explorations. One exploration means expanding a parent into 128 possible
             children. The number of calls to inverse_dct is then 128 * max_iter
-            upper_bound: theoretical upper bound between the starting block and the solution, depends on the pipeline.
             shared_dict: a manager shared memory to track tasks status
             task_id: an integer to track the task status
             verbose: if true, the function will use the task_id and the shared_dict to communicate advancement
@@ -74,10 +72,7 @@ class Block:
         open_set = {start.data.tobytes().__hash__()}
 
         s = c * n * m  # (channel * row * column)
-        if upper_bound is not None:
-            mask = np.ravel(pipeline.pipelines[0].quant_tbl.astype(float) <= upper_bound)
-        else:
-            mask = np.ones(s).astype(bool)  # True for every changes
+        mask = np.ravel(pipeline.pipelines[0].quant_tbl.astype(float) <= pipeline.upper_bound)
         n_changes = np.sum(mask)
         changes = np.stack([np.eye(s, dtype=np.int8)[mask],
                             -np.eye(s, dtype=np.int8)[mask]]).reshape(2 * n_changes, c, n, m)
@@ -99,10 +94,9 @@ class Block:
                 children_hash[i] = child.data.tobytes().__hash__()
                 not_ignored[i] = children_hash[i] not in open_set
             children.flags.writeable = True
-            if upper_bound is not None:
             distance = np.abs((children - float_start) * pipeline.pipelines[0].quant_tbl)
             norm_distance = np.linalg.norm(distance, axis=(2, 3), ord='fro')
-                not_ignored = not_ignored & np.all(norm_distance <= np.ravel(upper_bound), axis=-1)
+            not_ignored = not_ignored & np.all(norm_distance <= np.ravel(pipeline.upper_bound), axis=-1)
 
             if np.any(not_ignored):
                 transformed_children = pipeline.forward(children[not_ignored])

antecedent/pipeline.py

@@ -7,6 +7,7 @@ import numpy as np
 
 class JPEGPipeline:
     def __init__(self, name, quality, target_is_dct, grayscale, return_int=True, return_rounded=True):
+        self.upper_bound = None
         self.name = name
         self.quality = quality
         self.target_is_dct = target_is_dct
@@ -14,6 +15,7 @@ class JPEGPipeline:
         self.quant_tbl = None
         self.return_rounded = return_rounded
         self.define_quant_table()
+        self.define_upper_bound()
         if grayscale and target_is_dct and return_int:
             self.forward = self.pxl_to_dct
             self.backward = self.dct_to_pxl
@@ -51,6 +53,15 @@ class JPEGPipeline:
             h = (h * 281 ^ ord(ch) * 997) & 0xFFFFFFFF
         return h
 
+    def define_upper_bound(self):
+        self.upper_bound = np.zeros((self.quant_tbl.shape[0])).reshape(-1, 1, 1)
+        if self.target_is_dct:
+            self.upper_bound += np.linalg.norm(self.quant_tbl, axis=(1, 2), ord='fro', keepdims=True)
+        else:
+            self.upper_bound += 4
+        if not self.grayscale:
+            self.upper_bound += 4
+
     @classmethod
     def is_named(cls, name):
         return False
@@ -144,10 +155,12 @@ class IslowPipeline(JPEGPipeline):
         return jpeg_idct_islow(blocks, self.quant_tbl[:blocks.shape[1]])
 
     def pxl_to_dct(self, blocks):
-        return quantize_islow_fdct(jpeg_fdct_islow(blocks), quant_tbl=self.quant_tbl[:blocks.shape[1]], return_int=self.return_rounded)
+        return quantize_islow_fdct(jpeg_fdct_islow(blocks), quant_tbl=self.quant_tbl[:blocks.shape[1]],
+                                   return_int=self.return_rounded)
 
     def rgb_pxl_to_ycc_dct(self, blocks):
-        return quantize_islow_fdct(jpeg_fdct_islow(rgb_to_ycc(blocks)), quant_tbl=self.quant_tbl, return_int=self.return_rounded)
+        return quantize_islow_fdct(jpeg_fdct_islow(rgb_to_ycc(blocks)), quant_tbl=self.quant_tbl,
+                                   return_int=self.return_rounded)
 
     def dct_to_pxl(self, blocks):
         return jpeg_idct_islow(blocks, self.quant_tbl[:blocks.shape[1]])
@@ -191,8 +204,11 @@ class FloatPipeline(JPEGPipeline):
 class ComposedPipeline:
     def __init__(self, pipelines, backward_naive=True):
         self.pipelines = pipelines
+        self.upper_bound = np.sum([pipe.upper_bound for pipe in pipelines], axis=0)
         if backward_naive:
-            self.backward_pipes = [NaivePipeline(pipe.quality, pipe.target_is_dct, pipe.grayscale, return_int=False, return_rounded=False) for pipe in pipelines]
+            self.backward_pipes = [
+                NaivePipeline(pipe.quality, pipe.target_is_dct, pipe.grayscale, return_int=False, return_rounded=False)
+                for pipe in pipelines]
         else:
             self.backward_pipes = self.pipelines
         self.n = len(pipelines)
@@ -235,7 +251,9 @@ def create_pipeline(name, quality, grayscale, target_is_dct, backward_naive=True
             compression = not target_is_dct
         for subclass in JPEGPipeline.__subclasses__():
             if subclass.is_named(name[i].lower()):
-                pipelines.append(subclass(quality=quality[i], target_is_dct=compression, grayscale=grayscale, return_int=return_int, return_rounded=return_rounded))
+                pipelines.append(
+                    subclass(quality=quality[i], target_is_dct=compression, grayscale=grayscale, return_int=return_int,
+                             return_rounded=return_rounded))
                 found = True
                 break
         if found: