Browse Source

refactor: Simplify image rotation caching by consolidating image and text bounding box mapping into a single dictionary structure

zhch158_admin 19 giờ trước cách đây
mục cha
commit
01443814e4
1 tập tin đã thay đổi với 6 bổ sung9 xóa
  1. 6 9
      ocr_validator_layout.py

+ 6 - 9
ocr_validator_layout.py

@@ -86,10 +86,6 @@ class OCRLayoutManager:
                 if isinstance(item, dict) and 'rotation_angle' in item:
                     return item['rotation_angle']
         
-        # 如果没有预设角度,尝试自动检测
-        if hasattr(self, 'rotated_angle'):
-            return self.rotated_angle
-        
         return 0.0
     
     def load_and_rotate_image(self, image_path: str) -> Optional[Image.Image]:
@@ -102,7 +98,8 @@ class OCRLayoutManager:
         cache_key = f"{image_path}_{rotation_angle}"
         
         if cache_key in self._rotated_image_cache:
-            return self._rotated_image_cache[cache_key]
+            self.validator.text_bbox_mapping = self._rotated_image_cache[cache_key]['text_bbox_mapping']
+            return self._rotated_image_cache[cache_key]['image']
         
         try:
             image = Image.open(image_path)
@@ -132,13 +129,13 @@ class OCRLayoutManager:
                     
                     if self.rotated_angle == 0:
                         # 坐标不需要变换,因为JSON中已经是正确的坐标
-                        self._rotated_image_cache[cache_key] = rotated_image
+                        self._rotated_image_cache[cache_key] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                         self._manage_cache_size()
                         return rotated_image
 
                     image = rotated_image  # 继续使用旋转后的图像进行后续处理
                 
-                # Dots OCR: 需要同时旋转图像和坐标
+                # VLM: 需要同时旋转图像和坐标
                 # 收集所有bbox坐标
                 all_bboxes = []
                 text_to_bbox_map = {}  # 记录文本到bbox索引的映射
@@ -164,13 +161,13 @@ class OCRLayoutManager:
                             self.validator.text_bbox_mapping[text][i]['bbox'] = rotated_bboxes[bbox_idx]
                 
                 # 缓存结果
-                self._rotated_image_cache[cache_key] = rotated_image
+                self._rotated_image_cache[cache_key] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                 self._manage_cache_size()
                 return rotated_image
                     
             else:
                 # 无需旋转,直接缓存原图
-                self._rotated_image_cache[cache_key] = image
+                self._rotated_image_cache[cache_key] = {'image': image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                 self._manage_cache_size()  # 检查并管理缓存大小
                 return image