def apply_augmentation()

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

• 23 lines of code
• 5 McCabe index (conditional complexity)

    def apply_augmentation(self, raw_frame: np.ndarray, **kwargs) -> np.ndarray:
        """
        Adds shape distracts (in various colors 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
        ), "VideoDistractorByShapes only accepts RGB images"
        height, width = raw_frame.shape[:2]
        distract_frame = raw_frame.copy()

        for i in range(self.num_shapes):
            shape_type = next(self.shape_type)
            color = next(self.colors)
            thickness = next(self.thickness)
            fraction_x, fraction_y = self.get_origins(i)
            x = int(fraction_x * width)
            y = int(fraction_y * height)

            if shape_type == "circle":
                smaller_side = min(height, width)
                radius = int(next(self.radius) * smaller_side)
                cv2.circle(distract_frame, (x, y), radius, color, thickness)

            if shape_type == "rectangle":
                fraction_x, fraction_y = self.get_origins(i)
                x_2 = int(fraction_x * width)
                y_2 = int(fraction_y * height)
                cv2.rectangle(distract_frame, (x, y), (x_2, y_2), color, thickness)

        return distract_frame