def apply_augmentation()

in augly/video/augmenters/cv2/text.py [0:0]

• 41 lines of code
• 7 McCabe index (conditional complexity)

    def apply_augmentation(self, raw_frame: np.ndarray, **kwargs) -> np.ndarray:
        """
        Adds text distracts (in various colors, fonts, and positions) to each frame

        @param raw_frame: raw, single RGB/Gray frame

        @returns: the augumented frame
        """
        assert (raw_frame.ndim == 3) and (
            raw_frame.shape[2] == 3
        ), "VideoDistractorByText only accepts RGB images"
        height, width = raw_frame.shape[:2]

        text = next(self.texts)
        font, chars = next(self.fonts)  # pyre-ignore
        fontscale = next(self.fontscales)
        color = next(self.colors)
        thickness = next(self.thickness)
        fraction_x, fraction_y = self.get_origins(0)
        x = int(fraction_x * width)
        y = int(fraction_y * height)
        n = len(chars)
        text_str = "".join([chars[int(c * n)] for c in text])

        distract_frame = raw_frame.copy()
        if isinstance(font, str):
            with pathmgr.open(font, "rb") as f:
                # pyre-fixme[6]: Expected `Union[None,
                #  _typeshed.SupportsRead[bytes], bytes, str]` for 1st param but got
                #  `Union[typing.IO[bytes], typing.IO[str]]`.
                font = ImageFont.truetype(f, int(fontscale * 100))
        if isinstance(
            font,
            (ImageFont.ImageFont, ImageFont.FreeTypeFont, ImageFont.TransposedFont),
        ):
            # To use an ImageFont, we need to convert into PIL
            distract_frame_rgb = cv2.cvtColor(distract_frame, cv2.COLOR_BGR2RGB)
            distract_frame_pil = Image.fromarray(distract_frame_rgb)
            # pyre-fixme[6]: Expected `Optional[ImageFont._Font]` for 3rd param but
            #  got `Union[ImageFont.FreeTypeFont, ImageFont.ImageFont,
            #  ImageFont.TransposedFont]`.
            ImageDraw.Draw(distract_frame_pil).text((x, y), text_str, font=font)
            distract_frame = cv2.cvtColor(
                np.array(distract_frame_pil), cv2.COLOR_RGB2BGR
            )
        else:
            cv2.putText(
                distract_frame,
                text_str,
                (x, y),
                font,
                fontscale,
                color,
                thickness,
                cv2.LINE_AA,
            )
        return distract_frame