Esempio n. 1
0
    def apply(self, text_bboxes: List[BBox], img_bboxes: List[BBox],) -> List[BBox]:
        avg_height = sum([it.height for it in img_bboxes]) / len(img_bboxes)

        for i in range(0, len(img_bboxes) - 1):
            h_spacing_scale = np.random.uniform(*self.h_spacing)
            h_spacing = int(avg_height * h_spacing_scale)
            img_bboxes[i].right += h_spacing

        merged_bbox = BBox.from_bboxes(img_bboxes)

        img_bboxes[0].offset_(img_bboxes[0].left_cnt, merged_bbox.left_cnt)

        for i in range(1, len(img_bboxes)):
            img_bboxes[i].offset_(img_bboxes[i].left_cnt, img_bboxes[i - 1].right_cnt)

        return img_bboxes
    def gen_multi_corpus(self) -> Tuple[PILImage, str]:
        font_texts: List[FontText] = [it.sample() for it in self.corpus]

        bg = self.bg_manager.get_bg()

        text_color = None
        if self.cfg.text_color_cfg is not None:
            text_color = self.cfg.text_color_cfg.get_color(bg)

        text_masks, text_bboxes = [], []
        for i in range(len(font_texts)):
            font_text = font_texts[i]

            if text_color is None:
                _text_color = self.corpus[i].cfg.text_color_cfg.get_color(bg)
            else:
                _text_color = text_color
            text_mask = draw_text_on_bg(
                font_text,
                _text_color,
                char_spacing=self.corpus[i].cfg.char_spacing)

            text_bbox = BBox.from_size(text_mask.size)
            if self.cfg.corpus_effects is not None:
                effects = self.cfg.corpus_effects[i]
                if effects is not None:
                    text_mask, text_bbox = effects.apply_effects(
                        text_mask, text_bbox)
            text_masks.append(text_mask)
            text_bboxes.append(text_bbox)

        text_mask_bboxes, merged_text = self.layout(
            font_texts,
            [it.copy() for it in text_bboxes],
            [BBox.from_size(it.size) for it in text_masks],
        )
        if len(text_mask_bboxes) != len(text_bboxes):
            raise PanicError(
                "points and text_bboxes should have same length after layout output"
            )

        merged_bbox = BBox.from_bboxes(text_mask_bboxes)
        merged_text_mask = transparent_img(merged_bbox.size)
        for text_mask, bbox in zip(text_masks, text_mask_bboxes):
            merged_text_mask.paste(text_mask, bbox.left_top)

        if self.cfg.perspective_transform is not None:
            transformer = PerspectiveTransform(self.cfg.perspective_transform)
            # TODO: refactor this, now we must call get_transformed_size to call gen_warp_matrix
            _ = transformer.get_transformed_size(merged_text_mask.size)

            (
                transformed_text_mask,
                transformed_text_pnts,
            ) = transformer.do_warp_perspective(merged_text_mask)
        else:
            transformed_text_mask = merged_text_mask

        if self.cfg.layout_effects is not None:
            transformed_text_mask, _ = self.cfg.layout_effects.apply_effects(
                transformed_text_mask,
                BBox.from_size(transformed_text_mask.size))

        img = self.paste_text_mask_on_bg(bg, transformed_text_mask)

        return img, merged_text