Skip to content

Commit 77d8138

Browse files
committed
Add the function of concatenating to crops after detection.
1 parent bfb030d commit 77d8138

File tree

2 files changed

+28
-0
lines changed

2 files changed

+28
-0
lines changed

deploy/py_infer/src/infer_args.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,9 @@ def get_args():
119119
"--show_log", type=str2bool, default=False, required=False, help="Whether show log when inferring."
120120
)
121121
parser.add_argument("--save_log_dir", type=str, required=False, help="Log saving dir.")
122+
parser.add_argument(
123+
"--is_concat", type=str2bool, default=False, help="Whether to concatenate crops after the detection."
124+
)
122125

123126
args = parser.parse_args()
124127
setup_logger(args)

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import cv2
12
import numpy as np
23

34
from ....data_process.utils import cv_utils
@@ -16,6 +17,28 @@ def init_self_args(self):
1617
self.text_detector.init(preprocess=False, model=False, postprocess=True)
1718
super().init_self_args()
1819

20+
def concat_crops(self, crops: list):
21+
"""
22+
Concatenates the list of cropped images horizontally after resizing them to have the same height.
23+
24+
Args:
25+
crops (list): A list of cropped images represented as numpy arrays.
26+
27+
Returns:
28+
numpy.ndarray: A horizontally concatenated image array.
29+
"""
30+
max_height = max(crop.shape[0] for crop in crops)
31+
resized_crops = []
32+
for crop in crops:
33+
h, w, c = crop.shape
34+
new_h = max_height
35+
new_w = int((w / h) * new_h)
36+
37+
resized_img = cv2.resize(crop, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
38+
resized_crops.append(resized_img)
39+
crops = np.concatenate(resized_crops, axis=1)
40+
return crops
41+
1942
def process(self, input_data):
2043
if input_data.skip:
2144
self.send_to_next_module(input_data)
@@ -39,6 +62,8 @@ def process(self, input_data):
3962
for box in infer_res_list:
4063
sub_image = cv_utils.crop_box_from_image(image, np.array(box))
4164
sub_image_list.append(sub_image)
65+
if self.is_concat:
66+
sub_image_list = [self.concat_crops(sub_image_list)]
4267
input_data.sub_image_list = sub_image_list
4368

4469
input_data.data = None

0 commit comments

Comments
 (0)