chore: update files

Author: SWHL

demo.py

@@ -13,5 +13,5 @@
 table_engine = RapidTable(input_args)
 
 img_list = list(Path("images").iterdir())
-results = table_engine(img_path, batch_size=3)
+results = table_engine(img_list, batch_size=3)
 results.vis(save_dir="outputs", save_name="vis", indexes=(0, 1, 2))

rapid_table/main.py

@@ -126,7 +126,14 @@ def get_ocr_results(
         batch_dt_boxes, batch_rec_res = [], []
 
         if ocr_results is not None:
-            for img, ocr_result in zip(imgs, ocr_results[start_i:end_i]):
+            ocr_results_batch = ocr_results[start_i:end_i]
+            if len(ocr_results_batch) != len(imgs):
+                raise ValueError(
+                    f"Batch size mismatch: {len(imgs)} images but {len(ocr_results_batch)} OCR results "
+                    f"(indices {start_i}:{end_i})."
+                )
+
+            for img, ocr_result in zip(imgs, ocr_results_batch):
                 img_h, img_w = img.shape[:2]
                 dt_boxes, rec_res = format_ocr_results(ocr_result, img_h, img_w)
                 batch_dt_boxes.append(dt_boxes)

rapid_table/table_structure/unitable/unitable_modules.py

@@ -748,7 +748,7 @@ def forward(self, memory: Tensor, tgt: Tensor) -> Tensor:
                 logits = layer(logits, memory, input_pos=input_pos, tgt_mask=tgt_mask)
 
         logits = self.generator(logits)[:, -1, :]
-        total = set(list(range(logits.shape[-1])))
+        total = set(range(logits.shape[-1]))
         black_list = list(total.difference(set(self.token_white_list)))
         logits[..., black_list] = -1e9
         probs = F.softmax(logits, dim=-1)