Add the function of concatenating to crops after detection.

Author: Bourn3z

deploy/py_infer/src/infer_args.py

@@ -119,6 +119,9 @@ def get_args():
         "--show_log", type=str2bool, default=False, required=False, help="Whether show log when inferring."
     )
     parser.add_argument("--save_log_dir", type=str, required=False, help="Log saving dir.")
+    parser.add_argument(
+        "--is_concat", type=str2bool, default=False, help="Whether to concatenate crops after the detection."
+    )
 
     args = parser.parse_args()
     setup_logger(args)

deploy/py_infer/src/parallel/module/detection/det_post_node.py

@@ -1,3 +1,4 @@
+import cv2
 import numpy as np
 
 from ....data_process.utils import cv_utils
@@ -10,19 +11,44 @@ def __init__(self, args, msg_queue):
         super(DetPostNode, self).__init__(args, msg_queue)
         self.text_detector = None
         self.task_type = self.args.task_type
+        self.is_concat = self.args.is_concat
 
     def init_self_args(self):
         self.text_detector = TextDetector(self.args)
         self.text_detector.init(preprocess=False, model=False, postprocess=True)
         super().init_self_args()
 
+    def concat_crops(self, crops: list):
+        """
+        Concatenates the list of cropped images horizontally after resizing them to have the same height.
+
+        Args:
+            crops (list): A list of cropped images represented as numpy arrays.
+
+        Returns:
+            numpy.ndarray: A horizontally concatenated image array.
+        """
+        max_height = max(crop.shape[0] for crop in crops)
+        resized_crops = []
+        for crop in crops:
+            h, w, c = crop.shape
+            new_h = max_height
+            new_w = int((w / h) * new_h)
+
+            resized_img = cv2.resize(crop, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+            resized_crops.append(resized_img)
+        crops_concated = np.concatenate(resized_crops, axis=1)
+        return crops_concated
+
     def process(self, input_data):
         if input_data.skip:
             self.send_to_next_module(input_data)
             return
 
         data = input_data.data
         boxes = self.text_detector.postprocess(data["pred"], data["shape_list"])
+        if self.is_concat:
+            boxes = sorted(boxes, key=lambda points: (points[0][1], points[0][0]))
 
         infer_res_list = []
         for box in boxes:
@@ -39,6 +65,8 @@ def process(self, input_data):
             for box in infer_res_list:
                 sub_image = cv_utils.crop_box_from_image(image, np.array(box))
                 sub_image_list.append(sub_image)
+            if self.is_concat:
+                sub_image_list = [self.concat_crops(sub_image_list)]
             input_data.sub_image_list = sub_image_list
 
         input_data.data = None

deploy/py_infer/src/parallel/module/recognition/rec_post_node.py

@@ -7,6 +7,7 @@ def __init__(self, args, msg_queue):
         super(RecPostNode, self).__init__(args, msg_queue)
         self.text_recognizer = None
         self.task_type = self.args.task_type
+        self.is_concat = self.args.is_concat
 
     def init_self_args(self):
         self.text_recognizer = TextRecognizer(self.args)
@@ -28,9 +29,13 @@ def process(self, input_data):
         else:
             texts = output["texts"]
             confs = output["confs"]
-            for result, text, conf in zip(input_data.infer_result, texts, confs):
-                result.append(text)
-                result.append(conf)
+            for i, result in enumerate(input_data.infer_result):
+                if self.is_concat:
+                    result.append(texts[0])
+                    result.append(confs[0])
+                else:
+                    result.append(texts[i])
+                    result.append(confs[i])
 
         input_data.data = None