fix benchmark (#4404)

paddlex/inference/pipelines/_parallel.py

@@ -44,7 +44,7 @@ class MultiDeviceSimpleInferenceExecutor(object):
                 input_future_pairs = []
                 for pipeline in self._pipelines:
                     try:
-                        input_batch = next(input_batches)
+                        input_batch = list(input_batches)[0]
                     except StopIteration:
                         out_of_data = True
                         break

paddlex/inference/pipelines/formula_recognition/pipeline.py

@@ -268,7 +268,7 @@ class _FormulaRecognitionPipeline(BasePipeline):
                     layout_det_results = []
                     for _ in doc_preprocessor_images:
                         try:
-                            layout_det_res = next(external_layout_det_results)
+                            layout_det_res = list(external_layout_det_results)[0]
                         except StopIteration:
                             raise ValueError("No more layout det results")
                         layout_det_results.append(layout_det_res)

paddlex/inference/pipelines/keypoint_detection/pipeline.py

@@ -128,7 +128,7 @@ class _KeypointDetectionPipeline(BasePipeline):
             single_img_res = {"input_path": img_path, "input_img": ori_img, "boxes": []}
             for box in det_res["boxes"]:
                 center, scale = self._box_xyxy2cs(box["coordinate"])
-                kpt_res = next(
+                kpt_res = list(
                     self.kpt_model(
                         {
                             "img": ori_img,
@@ -136,7 +136,7 @@ class _KeypointDetectionPipeline(BasePipeline):
                             "scale": scale,
                         }
                     )
-                )
+                )[0]
                 single_img_res["boxes"].append(
                     {
                         "coordinate": box["coordinate"],

paddlex/inference/pipelines/layout_parsing/layout_objects.py

@@ -181,7 +181,7 @@ class TextLine(object):
                             int(bbox[1]) : int(bbox[3]),
                             int(bbox[0]) : int(bbox[2]),
                         ]
-                        crop_img_rec_res = next(text_rec_model([crop_img]))
+                        crop_img_rec_res = list(text_rec_model([crop_img]))[0]
                         crop_img_rec_score = crop_img_rec_res["rec_score"]
                         crop_img_rec_text = crop_img_rec_res["rec_text"]
                         span.text = crop_img_rec_text

paddlex/inference/pipelines/layout_parsing/pipeline.py

@@ -256,7 +256,7 @@ class _LayoutParsingPipeline(BasePipeline):
                     box = layout_parsing_res[idx]["block_bbox"]
                     x1, y1, x2, y2 = [int(i) for i in box]
                     wht_im[y1:y2, x1:x2, :] = image[y1:y2, x1:x2, :]
-                    sub_ocr_res = next(
+                    sub_ocr_res = list(
                         self.general_ocr_pipeline(
                             wht_im,
                             text_det_limit_side_len=text_det_limit_side_len,
@@ -266,7 +266,7 @@ class _LayoutParsingPipeline(BasePipeline):
                             text_det_unclip_ratio=text_det_unclip_ratio,
                             text_rec_score_thresh=text_rec_score_thresh,
                         )
-                    )
+                    )[0]
                     layout_parsing_res[idx]["block_content"] = "\n".join(
                         sub_ocr_res["rec_texts"]
                     )
@@ -443,19 +443,19 @@ class _LayoutParsingPipeline(BasePipeline):
             image_array = self.img_reader(batch_data.instances)[0]
 
             if model_settings["use_doc_preprocessor"]:
-                doc_preprocessor_res = next(
+                doc_preprocessor_res = list(
                     self.doc_preprocessor_pipeline(
                         image_array,
                         use_doc_orientation_classify=use_doc_orientation_classify,
                         use_doc_unwarping=use_doc_unwarping,
                     )
-                )
+                )[0]
             else:
                 doc_preprocessor_res = {"output_img": image_array}
 
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
 
-            layout_det_res = next(
+            layout_det_res = list(
                 self.layout_det_model(
                     doc_preprocessor_image,
                     threshold=layout_threshold,
@@ -463,9 +463,9 @@ class _LayoutParsingPipeline(BasePipeline):
                     layout_unclip_ratio=layout_unclip_ratio,
                     layout_merge_bboxes_mode=layout_merge_bboxes_mode,
                 )
-            )
+            )[0]
 
-            overall_ocr_res = next(
+            overall_ocr_res = list(
                 self.general_ocr_pipeline(
                     doc_preprocessor_image,
                     use_textline_orientation=use_textline_orientation,
@@ -476,10 +476,10 @@ class _LayoutParsingPipeline(BasePipeline):
                     text_det_unclip_ratio=text_det_unclip_ratio,
                     text_rec_score_thresh=text_rec_score_thresh,
                 )
-            )
+            )[0]
 
             if model_settings["use_table_recognition"]:
-                table_res_all = next(
+                table_res_all = list(
                     self.table_recognition_pipeline(
                         doc_preprocessor_image,
                         use_doc_orientation_classify=False,
@@ -489,13 +489,13 @@ class _LayoutParsingPipeline(BasePipeline):
                         overall_ocr_res=overall_ocr_res,
                         layout_det_res=layout_det_res,
                     )
-                )
+                )[0]
                 table_res_list = table_res_all["table_res_list"]
             else:
                 table_res_list = []
 
             if model_settings["use_seal_recognition"]:
-                seal_res_all = next(
+                seal_res_all = list(
                     self.seal_recognition_pipeline(
                         doc_preprocessor_image,
                         use_doc_orientation_classify=False,
@@ -509,13 +509,13 @@ class _LayoutParsingPipeline(BasePipeline):
                         seal_det_unclip_ratio=seal_det_unclip_ratio,
                         seal_rec_score_thresh=seal_rec_score_thresh,
                     )
-                )
+                )[0]
                 seal_res_list = seal_res_all["seal_res_list"]
             else:
                 seal_res_list = []
 
             if model_settings["use_formula_recognition"]:
-                formula_res_all = next(
+                formula_res_all = list(
                     self.formula_recognition_pipeline(
                         doc_preprocessor_image,
                         use_layout_detection=False,
@@ -523,7 +523,7 @@ class _LayoutParsingPipeline(BasePipeline):
                         use_doc_unwarping=False,
                         layout_det_res=layout_det_res,
                     )
-                )
+                )[0]
                 formula_res_list = formula_res_all["formula_res_list"]
             else:
                 formula_res_list = []

paddlex/inference/pipelines/layout_parsing/pipeline_v2.py

@@ -451,7 +451,7 @@ class _LayoutParsingPipelineV2(BasePipeline):
                 crop_box = layout_det_res["boxes"][layout_box_idx]["coordinate"]
                 x1, y1, x2, y2 = [int(i) for i in crop_box]
                 crop_img = np.array(image)[y1:y2, x1:x2]
-                crop_img_rec_res = next(text_rec_model([crop_img]))
+                crop_img_rec_res = list(text_rec_model([crop_img]))[0]
                 crop_img_dt_poly = get_bbox_intersection(
                     crop_box, crop_box, return_format="poly"
                 )

paddlex/inference/pipelines/seal_recognition/pipeline.py

@@ -259,7 +259,7 @@ class _SealRecognitionPipeline(BasePipeline):
                     layout_det_results = []
                     for _ in doc_preprocessor_images:
                         try:
-                            layout_det_res = next(external_layout_det_results)
+                            layout_det_res = list(external_layout_det_results)[0]
                         except StopIteration:
                             raise ValueError("No more layout det results")
                         layout_det_results.append(layout_det_res)

paddlex/inference/pipelines/table_recognition/pipeline.py

@@ -218,13 +218,13 @@ class _TableRecognitionPipeline(BasePipeline):
         if input_params["use_doc_preprocessor"]:
             use_doc_orientation_classify = input_params["use_doc_orientation_classify"]
             use_doc_unwarping = input_params["use_doc_unwarping"]
-            doc_preprocessor_res = next(
+            doc_preprocessor_res = list(
                 self.doc_preprocessor_pipeline(
                     image_array,
                     use_doc_orientation_classify=use_doc_orientation_classify,
                     use_doc_unwarping=use_doc_unwarping,
                 )
-            )
+            )[0]
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
         else:
             doc_preprocessor_res = {}
@@ -252,7 +252,7 @@ class _TableRecognitionPipeline(BasePipeline):
             # Extract and round up the coordinates of the bounding box.
             x1, y1, x2, y2 = [math.ceil(k) for k in cells_bboxes[i]]
             # Perform OCR on the defined region of the image and get the recognized text.
-            rec_te = next(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))
+            rec_te = list(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))[0]
             # Concatenate the texts and append them to the texts_list.
             texts_list.append("".join(rec_te["rec_texts"]))
         # Return the list of recognized texts from each cell.
@@ -279,7 +279,7 @@ class _TableRecognitionPipeline(BasePipeline):
             # Extract and round up the coordinates of the bounding box.
             x1, y1, x2, y2 = [math.ceil(k) for k in cells_bboxes[i]]
             # Perform OCR on the defined region of the image and get the recognized text.
-            rec_te = next(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))
+            rec_te = list(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))[0]
             # Concatenate the texts and append them to the texts_list.
             texts_list.append("".join(rec_te["rec_texts"]))
         # Return the list of recognized texts from each cell.
@@ -308,7 +308,7 @@ class _TableRecognitionPipeline(BasePipeline):
         Returns:
             SingleTableRecognitionResult: single table recognition result.
         """
-        table_structure_pred = next(self.table_structure_model(image_array))
+        table_structure_pred = list(self.table_structure_model(image_array))[0]
         if use_ocr_results_with_table_cells == True:
             table_cells_result = table_structure_pred["bbox"]
             table_cells_result = [
@@ -395,20 +395,20 @@ class _TableRecognitionPipeline(BasePipeline):
             image_array = self.img_reader(batch_data.instances)[0]
 
             if model_settings["use_doc_preprocessor"]:
-                doc_preprocessor_res = next(
+                doc_preprocessor_res = list(
                     self.doc_preprocessor_pipeline(
                         image_array,
                         use_doc_orientation_classify=use_doc_orientation_classify,
                         use_doc_unwarping=use_doc_unwarping,
                     )
-                )
+                )[0]
             else:
                 doc_preprocessor_res = {"output_img": image_array}
 
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
 
             if model_settings["use_ocr_model"]:
-                overall_ocr_res = next(
+                overall_ocr_res = list(
                     self.general_ocr_pipeline(
                         doc_preprocessor_image,
                         text_det_limit_side_len=text_det_limit_side_len,
@@ -418,7 +418,7 @@ class _TableRecognitionPipeline(BasePipeline):
                         text_det_unclip_ratio=text_det_unclip_ratio,
                         text_rec_score_thresh=text_rec_score_thresh,
                     )
-                )
+                )[0]
             elif use_ocr_results_with_table_cells == True:
                 assert self.general_ocr_config_bak != None
                 self.general_ocr_pipeline = self.create_pipeline(
@@ -444,7 +444,9 @@ class _TableRecognitionPipeline(BasePipeline):
                 table_region_id += 1
             else:
                 if model_settings["use_layout_detection"]:
-                    layout_det_res = next(self.layout_det_model(doc_preprocessor_image))
+                    layout_det_res = list(
+                        self.layout_det_model(doc_preprocessor_image)
+                    )[0]
 
                 for box_info in layout_det_res["boxes"]:
                     if box_info["label"].lower() in ["table"]:

paddlex/inference/pipelines/table_recognition/pipeline_v2.py

@@ -267,13 +267,13 @@ class _TableRecognitionPipelineV2(BasePipeline):
         if input_params["use_doc_preprocessor"]:
             use_doc_orientation_classify = input_params["use_doc_orientation_classify"]
             use_doc_unwarping = input_params["use_doc_unwarping"]
-            doc_preprocessor_res = next(
+            doc_preprocessor_res = list(
                 self.doc_preprocessor_pipeline(
                     image_array,
                     use_doc_orientation_classify=use_doc_orientation_classify,
                     use_doc_unwarping=use_doc_unwarping,
                 )
-            )
+            )[0]
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
         else:
             doc_preprocessor_res = {}
@@ -686,11 +686,11 @@ class _TableRecognitionPipelineV2(BasePipeline):
                 for box in split_boxes:
                     x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
                     if y2 - y1 > 1 and x2 - x1 > 1:
-                        ocr_result = next(
+                        ocr_result = list(
                             self.general_ocr_pipeline.text_rec_model(
                                 ori_img[y1:y2, x1:x2, :]
                             )
-                        )
+                        )[0]
                         # Extract the recognized text from the OCR result
                         if "rec_text" in ocr_result:
                             result = ocr_result[
@@ -738,7 +738,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
             x1, y1, x2, y2 = [math.ceil(k) for k in cells_bboxes[i]]
             # Perform OCR on the defined region of the image and get the recognized text.
             if y2 - y1 > 1 and x2 - x1 > 1:
-                rec_te = next(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))
+                rec_te = list(self.general_ocr_pipeline(ori_img[y1:y2, x1:x2, :]))[0]
                 # Concatenate the texts and append them to the texts_list.
                 texts_list.append("".join(rec_te["rec_texts"]))
         # Return the list of recognized texts from each cell.
@@ -979,7 +979,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
             SingleTableRecognitionResult: single table recognition result.
         """
 
-        table_cls_pred = next(self.table_cls_model(image_array))
+        table_cls_pred = list(self.table_cls_model(image_array))[0]
         table_cls_result = self.extract_results(table_cls_pred, "cls")
         use_e2e_model = False
         cells_trans_to_html = False
@@ -988,33 +988,41 @@ class _TableRecognitionPipelineV2(BasePipeline):
             if use_wired_table_cells_trans_to_html == True:
                 cells_trans_to_html = True
             else:
-                table_structure_pred = next(self.wired_table_rec_model(image_array))
+                table_structure_pred = list(self.wired_table_rec_model(image_array))[0]
             if use_e2e_wired_table_rec_model == True:
                 use_e2e_model = True
                 if cells_trans_to_html == True:
-                    table_structure_pred = next(self.wired_table_rec_model(image_array))
+                    table_structure_pred = list(
+                        self.wired_table_rec_model(image_array)
+                    )[0]
             else:
-                table_cells_pred = next(
+                table_cells_pred = list(
                     self.wired_table_cells_detection_model(image_array, threshold=0.3)
-                )  # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
+                )[
+                    0
+                ]  # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
                 # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
         elif table_cls_result == "wireless_table":
             if use_wireless_table_cells_trans_to_html == True:
                 cells_trans_to_html = True
             else:
-                table_structure_pred = next(self.wireless_table_rec_model(image_array))
+                table_structure_pred = list(self.wireless_table_rec_model(image_array))[
+                    0
+                ]
             if use_e2e_wireless_table_rec_model == True:
                 use_e2e_model = True
                 if cells_trans_to_html == True:
-                    table_structure_pred = next(
+                    table_structure_pred = list(
                         self.wireless_table_rec_model(image_array)
-                    )
+                    )[0]
             else:
-                table_cells_pred = next(
+                table_cells_pred = list(
                     self.wireless_table_cells_detection_model(
                         image_array, threshold=0.3
                     )
-                )  # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
+                )[
+                    0
+                ]  # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
                 # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
 
         if use_e2e_model == False:
@@ -1172,20 +1180,20 @@ class _TableRecognitionPipelineV2(BasePipeline):
             image_array = self.img_reader(batch_data.instances)[0]
 
             if model_settings["use_doc_preprocessor"]:
-                doc_preprocessor_res = next(
+                doc_preprocessor_res = list(
                     self.doc_preprocessor_pipeline(
                         image_array,
                         use_doc_orientation_classify=use_doc_orientation_classify,
                         use_doc_unwarping=use_doc_unwarping,
                     )
-                )
+                )[0]
             else:
                 doc_preprocessor_res = {"output_img": image_array}
 
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
 
             if model_settings["use_ocr_model"]:
-                overall_ocr_res = next(
+                overall_ocr_res = list(
                     self.general_ocr_pipeline(
                         doc_preprocessor_image,
                         text_det_limit_side_len=text_det_limit_side_len,
@@ -1195,7 +1203,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
                         text_det_unclip_ratio=text_det_unclip_ratio,
                         text_rec_score_thresh=text_rec_score_thresh,
                     )
-                )
+                )[0]
             elif self.general_ocr_pipeline is None and (
                 (
                     use_ocr_results_with_table_cells == True
@@ -1218,9 +1226,9 @@ class _TableRecognitionPipelineV2(BasePipeline):
                 img_height, img_width = doc_preprocessor_image.shape[:2]
                 table_box = [0, 0, img_width - 1, img_height - 1]
                 if use_table_orientation_classify == True:
-                    table_angle = next(
+                    table_angle = list(
                         self.table_orientation_classify_model(doc_preprocessor_image)
-                    )["label_names"][0]
+                    )[0]["label_names"][0]
                 if table_angle == "90":
                     doc_preprocessor_image = np.rot90(doc_preprocessor_image, k=1)
                 elif table_angle == "180":
@@ -1228,7 +1236,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
                 elif table_angle == "270":
                     doc_preprocessor_image = np.rot90(doc_preprocessor_image, k=3)
                 if table_angle in ["90", "180", "270"]:
-                    overall_ocr_res = next(
+                    overall_ocr_res = list(
                         self.general_ocr_pipeline(
                             doc_preprocessor_image,
                             text_det_limit_side_len=text_det_limit_side_len,
@@ -1238,7 +1246,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
                             text_det_unclip_ratio=text_det_unclip_ratio,
                             text_rec_score_thresh=text_rec_score_thresh,
                         )
-                    )
+                    )[0]
                     tbx1, tby1, tbx2, tby2 = (
                         table_box[0],
                         table_box[1],
@@ -1282,7 +1290,9 @@ class _TableRecognitionPipelineV2(BasePipeline):
                 table_region_id += 1
             else:
                 if model_settings["use_layout_detection"]:
-                    layout_det_res = next(self.layout_det_model(doc_preprocessor_image))
+                    layout_det_res = list(
+                        self.layout_det_model(doc_preprocessor_image)
+                    )[0]
                 img_height, img_width = doc_preprocessor_image.shape[:2]
                 for box_info in layout_det_res["boxes"]:
                     if box_info["label"].lower() in ["table"]:
@@ -1293,11 +1303,11 @@ class _TableRecognitionPipelineV2(BasePipeline):
                         table_box = crop_img_info["box"]
                         if use_table_orientation_classify == True:
                             doc_preprocessor_image_copy = doc_preprocessor_image.copy()
-                            table_angle = next(
+                            table_angle = list(
                                 self.table_orientation_classify_model(
                                     crop_img_info["img"]
                                 )
-                            )["label_names"][0]
+                            )[0]["label_names"][0]
                         if table_angle == "90":
                             crop_img_info["img"] = np.rot90(crop_img_info["img"], k=1)
                             doc_preprocessor_image_copy = np.rot90(
@@ -1314,7 +1324,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
                                 doc_preprocessor_image_copy, k=3
                             )
                         if table_angle in ["90", "180", "270"]:
-                            overall_ocr_res = next(
+                            overall_ocr_res = list(
                                 self.general_ocr_pipeline(
                                     doc_preprocessor_image_copy,
                                     text_det_limit_side_len=text_det_limit_side_len,
@@ -1324,7 +1334,7 @@ class _TableRecognitionPipelineV2(BasePipeline):
                                     text_det_unclip_ratio=text_det_unclip_ratio,
                                     text_rec_score_thresh=text_rec_score_thresh,
                                 )
-                            )
+                            )[0]
                             tbx1, tby1, tbx2, tby2 = (
                                 table_box[0],
                                 table_box[1],

paddlex/inference/utils/benchmark.py

@@ -220,16 +220,17 @@ class Benchmark:
         def wrapper():
             global _level
 
-            while True:
-                tic = time.perf_counter()
-                try:
-                    item = next(generator)
-                except StopIteration:
-                    break
-                self._update(time.perf_counter() - tic, name)
-                yield item
-
-            _level -= 1
+            try:
+                while True:
+                    tic = time.perf_counter()
+                    try:
+                        item = next(generator)
+                    except StopIteration:
+                        break
+                    self._update(time.perf_counter() - tic, name)
+                    yield item
+            finally:
+                _level -= 1
 
         return wrapper()