doc_preprocess pipeline support to save multi-imgs result

paddlex/inference/pipelines_new/doc_preprocessor/result.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Dict
 import math
 import random
 from pathlib import Path
@@ -45,34 +46,14 @@ class DocPreprocessorResult(BaseCVResult):
             save_path = Path(save_path) / f"res_doc_preprocess_{img_id}.jpg"
         super().save_to_img(save_path, *args, **kwargs)
 
-    def _to_img(self) -> PIL.Image:
+    def _to_img(self) -> Dict[str, Image.Image]:
         """
         Generate an image combining the original, rotated, and unwarping images.
 
         Returns:
-            PIL.Image: A new image that displays the original, rotated, and unwarping images side by side.
+            Dict[Image.Image]: A new image that displays the rotated, and unwarping images.
         """
-        image = self["input_image"][:, :, ::-1]
-        angle = self["angle"]
-        rot_img = self["rot_img"][:, :, ::-1]
-        output_img = self["output_img"][:, :, ::-1]
-        h1, w1 = image.shape[0:2]
-        h2, w2 = rot_img.shape[0:2]
-        h3, w3 = output_img.shape[0:2]
-        h = max(max(h1, h2), h3)
-        img_show = Image.new("RGB", (w1 + w2 + w3, h + 25), (255, 255, 255))
-        img_show.paste(Image.fromarray(image), (0, 0, w1, h1))
-        img_show.paste(Image.fromarray(rot_img), (w1, 0, w1 + w2, h2))
-        img_show.paste(Image.fromarray(output_img), (w1 + w2, 0, w1 + w2 + w3, h3))
-
-        draw_text = ImageDraw.Draw(img_show)
-        txt_list = ["Original Image", "Rotated Image", "Unwarping Image"]
-        region_w_list = [w1, w2, w3]
-        beg_w_list = [0, w1, w1 + w2]
-        for tno in range(len(txt_list)):
-            txt = txt_list[tno]
-            font = create_font(txt, (region_w_list[tno], 20), PINGFANG_FONT_FILE_PATH)
-            draw_text.text(
-                [10 + beg_w_list[tno], h + 2], txt, fill=(0, 0, 0), font=font
-            )
-        return img_show
+        imgs = {"preprocessed_img": Image.fromarray(self["output_img"][:, :, ::-1])}
+        if self["rot_img"] is not None:
+            imgs["rotated_img"] = Image.fromarray(self["rot_img"][:, :, ::-1])
+        return imgs

paddlex/inference/pipelines_new/ocr/result.py

@@ -173,7 +173,7 @@ class OCRResult(BaseCVResult):
         img_id = self["img_id"]
 
         return {
-            f"doc_preprocessor_result_img_{img_id}": self["doc_preprocessor_res"].img,
+            **self["doc_preprocessor_res"].img,
             f"res_ocr_{img_id}": img_show,
         }