fix: fixed issue #112

SWHL · SWHL · commit 51bb0ecf66bd · 2025-07-16T09:22:13.000+08:00
diff --git a/rapid_table/main.py b/rapid_table/main.py
@@ -86,6 +86,7 @@ def __call__(
 
         dt_boxes, rec_res = self.get_ocr_results(img, ocr_results)
         pred_structures, cell_bboxes, logic_points = self.get_table_rec_results(img)
+
         pred_html = self.get_table_matcher(
             pred_structures, cell_bboxes, dt_boxes, rec_res
         )
diff --git a/rapid_table/table_structure/unitable/main.py b/rapid_table/table_structure/unitable/main.py
@@ -66,12 +66,6 @@ def __init__(self, cfg: Dict[str, Any]):
         )
 
         self.decoder = self.model.decoder
-        self.decoder.setup_caches(
-            max_batch_size=1,
-            max_seq_length=self.max_seq_len,
-            dtype=torch.float32,
-            device=self.device,
-        )
 
     @torch.inference_mode()
     def __call__(self, image: np.ndarray) -> Tuple[List[str], np.ndarray, float]:
@@ -80,6 +74,13 @@ def __call__(self, image: np.ndarray) -> Tuple[List[str], np.ndarray, float]:
         ori_h, ori_w = image.shape[:2]
         image = self.preprocess_img(image)
 
+        self.decoder.setup_caches(
+            max_batch_size=1,
+            max_seq_length=self.max_seq_len,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
         memory = self.encoder(image)
         context = self.loop_decode(self.context, self.eos_id_tensor, memory)
         bboxes, html_tokens = self.decode_tokens(context)

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,7 @@ def __call__(`
`86`	`86`
`87`	`87`	`dt_boxes, rec_res = self.get_ocr_results(img, ocr_results)`
`88`	`88`	`pred_structures, cell_bboxes, logic_points = self.get_table_rec_results(img)`
	`89`	`+`
`89`	`90`	`pred_html = self.get_table_matcher(`
`90`	`91`	`pred_structures, cell_bboxes, dt_boxes, rec_res`
`91`	`92`	`)`