fix: correct split search when there is no color conversion

antecedent/block.py

@@ -55,7 +55,7 @@ class Block:
         if np.allclose(pipeline.forward(start), target):
             send_logs(verbose, shared_dict, task_id, max_iter, max_iter)
             self.status[pipeline] = 1
-            self.antecedents[pipeline] = (start, 0)
+            self.antecedents[pipeline] = (start[0], 0)
             return self.antecedents[pipeline]
 
         queue = [(np.zeros(0), 0., start.astype(np.int16).copy())]

antecedent/image.py

@@ -41,6 +41,9 @@ class Image:
             self.is_jpeg = True
             self.img = jpeglib.read_dct(path)
             self.is_grayscale = not self.img.has_chrominance
+            if self.is_grayscale:
+                self.data = np.stack([self.img.Y], axis=2)  # shape (None, None, 3, 8, 8)
+            else:
                 self.data = np.stack([self.img.Y, self.img.Cb, self.img.Cr], axis=2)  # shape (None, None, 3, 8, 8)
             for i in range(self.data.shape[0]):
                 for j in range(self.data.shape[1]):
@@ -74,14 +77,14 @@ class Image:
         if self.rng is None:
             self.rng = np.random.RandomState()
 
-    def search_antecedent(self, max_iter, shared_dict=None, task_id=None, verbose=False):
+    def search_antecedent(self, max_iter, shared_dict=None, task_id=None, verbose=False, split_search_if_possible=True):
         self.filter_block()
         blocks = self.select_blocks()
-        if self.is_jpeg and self.pipeline.n == 1:
+        if self.is_jpeg and self.pipeline.n == 1 and split_search_if_possible:
             # Single pipeline from Pixel to DCT, we can search a pixel antecedent for each channel independently
             # A similar pipeline is created but without color conversion
-            names = [pipe.name for pipe in self.pipeline]
-            qualities = [pipe.quality for pipe in self.pipeline]
+            names = [pipe.name for pipe in self.pipeline.pipelines]
+            qualities = [pipe.quality for pipe in self.pipeline.pipelines]
             self.grayscale_pipeline = create_pipeline(names, qualities, grayscale=False, target_is_dct=True)
 
         for i, block in enumerate(blocks):
@@ -92,13 +95,13 @@ class Image:
                                         "completed": 0,
                                         "total": 0}
 
-            if self.is_jpeg and self.pipeline.n == 1:
+            if self.is_jpeg and self.pipeline.n == 1 and split_search_if_possible:
                 # Case where the channel are independent
                 antecedents = []
                 iterations = []
                 status = 1
-                for channel in block.shape[0]:
-                    single_channel_block = Block(channel, pos=block.pos)
+                for channel in block.value:
+                    single_channel_block = Block(np.array([channel]), pos=block.pos)
                     antecedent, iteration = single_channel_block.search_antecedent(self.grayscale_pipeline,
                                                                                    max_iter,
                                                                                    shared_dict,
@@ -149,7 +152,7 @@ class Image:
             return
         else:
             for _, block in self.block_collection.items():
-                spatial_block = self.pipeline.round_pixel(self.pipeline.inverse_dct(np.array([block.value])))
+                spatial_block = self.pipeline.backward(np.array([block.value]))
                 clipped = np.any(spatial_block >= 255) | np.any(spatial_block <= 0)
                 uniform = np.unique(spatial_block).size == 1
                 if self.avoid_clipped & clipped or self.avoid_uniform & uniform:
@@ -161,8 +164,9 @@ class Image:
             return self.selected_blocks
         n = min(int(len(self.block_collection) * self.selection_percentage / 100), len(self.block_collection))
         if self.method_name == 'random':
-            self.selected_blocks = self.rng.choice([b for _, b in self.block_collection.items() if not b.ignored],
-                                                   size=n,
+            potential_blocks = [b for _, b in self.block_collection.items() if not b.ignored]
+            self.selected_blocks = self.rng.choice(potential_blocks,
+                                                   size=min(len(potential_blocks), n),
                                                    replace=False)
             return self.selected_blocks
         else: