def apply(
self,
text_bboxes: List[BBox],
img_bboxes: List[BBox],
) -> List[BBox]:
assert (len(text_bboxes) == 2
), "ExtraTextLineLayout only support input two text bboxes"
main_text_mask_bbox, extra_text_mask_bbox = (
img_bboxes[0],
img_bboxes[1],
)
if prob(0.5):
# above extra line
extra_text_mask_bbox.offset_(
extra_text_mask_bbox.left_top,
(
main_text_mask_bbox.left_top[0],
main_text_mask_bbox.left_top[1] -
int(extra_text_mask_bbox.height * 0.8),
),
)
else:
# bottom extra line
extra_text_mask_bbox.offset_(extra_text_mask_bbox.left_top,
main_text_mask_bbox.left_bottom)
extra_text_mask_bbox.bottom -= int(extra_text_mask_bbox.height *
0.8)
if extra_text_mask_bbox.width > main_text_mask_bbox.width:
extra_text_mask_bbox.right -= (extra_text_mask_bbox.width -
main_text_mask_bbox.width)
return img_bboxes
def apply(
self,
text_bboxes: List[BBox],
img_bboxes: List[BBox],
) -> List[BBox]:
assert (len(text_bboxes) == 2
), "ExtraTextLineLayout only support input two text bboxes"
main_text_mask_bbox, extra_text_mask_bbox = (
img_bboxes[0],
img_bboxes[1],
)
if prob(1 - self.bottom_prob):
# above extra line
main_offset = int(main_text_mask_bbox.height *
random.uniform(0.3, 0.4))
main_text_mask_bbox.offset_(
main_text_mask_bbox.left_top,
(
main_text_mask_bbox.left_top[0],
main_text_mask_bbox.left_top[1] + main_offset,
),
)
extra_offset = int(random.uniform(0, main_offset // 2))
extra_text_mask_bbox.offset_(
extra_text_mask_bbox.left_bottom,
(
main_text_mask_bbox.left_top[0],
main_text_mask_bbox.left_top[1] - extra_offset,
),
)
else:
# bottom extra line
extra_text_mask_bbox.offset_(extra_text_mask_bbox.left_top,
main_text_mask_bbox.left_bottom)
extra_text_mask_bbox.bottom -= int(extra_text_mask_bbox.height *
random.uniform(0.65, 0.9))
if extra_text_mask_bbox.width > main_text_mask_bbox.width:
extra_text_mask_bbox.right -= (extra_text_mask_bbox.width -
main_text_mask_bbox.width)
return img_bboxes
def __call__(self, img, text_bbox):
if prob(self.p):
return self.apply(img, text_bbox)
return img, text_bbox