feat: add block iterations dict

README.md

@@ -50,7 +50,8 @@ for pos, block in img.block_collection.items():
         # 1 if an antecedent has been found
         # 0 otherwise (potentially incompatible)
         status = block.status[pipeline]
-        antecedent = block.antecedents[pipeline][0]
+        antecedent = block.antecedents[pipeline]
+        iteration = block.iterations[pipeline]
 ```
 
 If you have a single block:
@@ -64,7 +65,9 @@ array = np.random.randint(-1016, 1016, size=(8, 8))  # random DCT block
 block = Block(array)
 pipeline = create_pipeline('float', 90, grayscale=True, target_is_dct=True)  # Float pipeline with QF90 for grayscale blocks
 
-antecedent, iteration = block.search_antecedent(pipeline, 1000)
+block.search_antecedent(pipeline, 1000)
+antecedent = block.antecedents[pipeline]
+iteration = block.iterations[pipeline]
 if antecedent is None:
     print('No antecedent found, could be incompatible.')
 else:

antecedent/block.py

@@ -22,13 +22,15 @@ class Block:
         self.value = value
         self.antecedents = dict()
         self.status = dict()
+        self.iterations = dict()
         self.ignored = False
         self.pos = pos
 
     def ignore(self, pipeline):
         self.ignored = True
         self.status[pipeline] = 3
-        self.antecedents[pipeline] = (None, 0)
+        self.antecedents[pipeline] = None
+        self.iterations[pipeline] = 0
 
     def search_antecedent(self,
                           pipeline: ComposedPipeline,
@@ -63,7 +65,8 @@ class Block:
         if np.allclose(pipeline.forward(start), target):
             send_logs(verbose, shared_dict, task_id, max_iter, max_iter, True)
             self.status[pipeline] = 1
-            self.antecedents[pipeline] = (start, 0)
+            self.antecedents[pipeline] = start
+            self.iterations[pipeline] = 0
             return self.antecedents[pipeline]
 
         queue = [(np.zeros(0), 0., start.astype(np.int16).copy())]
@@ -110,8 +113,9 @@ class Block:
                 if error[error_idx[0]] == 0:  # check only the first element which is the smallest error
                     send_logs(verbose, shared_dict, task_id, max_iter, max_iter, True)
 
-                    self.antecedents[pipeline] = (children[not_ignored][error_idx[0]], iter_counter)
+                    self.antecedents[pipeline] = children[not_ignored][error_idx[0]]
                     self.status[pipeline] = 1
+                    self.iterations[pipeline] = iter_counter
 
                     return self.antecedents[pipeline]
 
@@ -124,11 +128,12 @@ class Block:
                         open_set.add(children_hash[not_ignored][idx])
 
         send_logs(verbose, shared_dict, task_id, max_iter, max_iter, True)
+        self.iterations[pipeline] = iter_counter + 1
         if iter_counter < max_iter:
-            self.antecedents[pipeline] = (False, iter_counter + 1)
+            self.antecedents[pipeline] = False
             self.status[pipeline] = -1
         else:
-            self.antecedents[pipeline] = (None, iter_counter + 1)
+            self.antecedents[pipeline] = None
             self.status[pipeline] = 0
         return self.antecedents[pipeline]
 

antecedent/image.py

@@ -99,18 +99,20 @@ class Image:
                 status = 1
                 for channel in block.shape[0]:
                     single_channel_block = Block(channel, pos=block.pos)
-                    antecedent, iteration = single_channel_block.search_antecedent(self.grayscale_pipeline,
+                    antecedent = single_channel_block.search_antecedent(self.grayscale_pipeline,
                                                                                    max_iter,
                                                                                    shared_dict,
                                                                                    task_id,
                                                                                    verbose)
+                    iteration = single_channel_block.iterations[self.grayscale_pipeline]
                     if antecedent is None:
                         antecedent = -np.ones((1, 8, 8))
                         status = 0
                     antecedents.append(antecedent)
                     iterations.append(iteration)
 
-                block.antecedents[self.pipeline] = (np.stack(antecedents, axis=0), np.sum(iterations))
+                block.antecedents[self.pipeline] = np.stack(antecedents, axis=0)
+                block.iterations[self.pipeline] = np.sum(iterations)
                 block.status[self.pipeline] = status
             else:
                 block.search_antecedent(self.pipeline, max_iter, shared_dict, task_id, verbose)

antecedent/run_on_dataset.py

@@ -121,13 +121,13 @@ def to_csv(path, image):
         for block in image.select_blocks():
             if pipeline in block.status:
                 status = block.status[pipeline]
-                antecedent = block.antecedents[pipeline][0]
+                antecedent = block.antecedents[pipeline]
+                iteration_heuristic = block.iterations[pipeline]
+                iteration_gurobi = 0
                 if antecedent is not None:
                     sol = str(list(np.ravel(antecedent).astype(int)))
                 else:
                     sol = None
-                iteration_heuristic = block.antecedents[pipeline][1]
-                iteration_gurobi = 0
 
                 writer.writerow(
                     [filename, str(block.pos), str(pipeline), str(status), str(iteration_heuristic),