Selaa lähdekoodia

perf(model): optimize batch analyze process

- Move batch model initialization outside the loop
- Collect page dimensions before analyzing- Update page info dictionary structure
- Add null dimensions for non-analyzed pages
myhloli 9 kuukautta sitten
vanhempi
commit
6753df8dd3
1 muutettua tiedostoa jossa 9 lisäystä ja 7 poistoa
  1. 9 7
      magic_pdf/model/doc_analyze_by_custom_model.py

+ 9 - 7
magic_pdf/model/doc_analyze_by_custom_model.py

@@ -157,6 +157,7 @@ def doc_analyze(
     )
 
     batch_analyze = False
+    batch_ratio = 1
     device = get_device()
 
     npu_support = False
@@ -181,7 +182,6 @@ def doc_analyze(
                 batch_ratio = 2
 
             logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}')
-            batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)
             batch_analyze = True
 
     model_json = []
@@ -190,24 +190,26 @@ def doc_analyze(
     if batch_analyze:
         # batch analyze
         images = []
+        page_wh_list = []
         for index in range(len(dataset)):
             if start_page_id <= index <= end_page_id:
                 page_data = dataset.get_page(index)
                 img_dict = page_data.get_image()
                 images.append(img_dict['img'])
+                page_wh_list.append((img_dict['width'], img_dict['height']))
+        batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)
         analyze_result = batch_model(images)
 
         for index in range(len(dataset)):
-            page_data = dataset.get_page(index)
-            img_dict = page_data.get_image()
-            page_width = img_dict['width']
-            page_height = img_dict['height']
             if start_page_id <= index <= end_page_id:
                 result = analyze_result.pop(0)
+                page_width, page_height = page_wh_list.pop(0)
             else:
                 result = []
+                page_height = 0
+                page_width = 0
 
-            page_info = {'page_no': index, 'height': page_height, 'width': page_width}
+            page_info = {'page_no': index, 'width': page_width, 'height': page_height}
             page_dict = {'layout_dets': result, 'page_info': page_info}
             model_json.append(page_dict)
 
@@ -227,7 +229,7 @@ def doc_analyze(
             else:
                 result = []
 
-            page_info = {'page_no': index, 'height': page_height, 'width': page_width}
+            page_info = {'page_no': index, 'width': page_width, 'height': page_height}
             page_dict = {'layout_dets': result, 'page_info': page_info}
             model_json.append(page_dict)