angrybeanie_wagtail/env/lib/python3.12/site-packages/willow/plugins/opencv.py

import os

from willow.image import Image, RGBImageBuffer


def _cv2():
    try:
        import cv2
    except ImportError:
        from cv import cv2
    return cv2


def _numpy():
    import numpy

    return numpy


class BaseOpenCVImage(Image):
    def __init__(self, image, size):
        self.image = image
        self.size = size

    @classmethod
    def check(cls):
        _cv2()

    @Image.operation
    def get_size(self):
        return self.size

    @Image.operation
    def get_frame_count(self):
        # Animation is not supported by OpenCV
        return 1

    @Image.operation
    def has_alpha(self):
        # Alpha is not supported by OpenCV
        return False

    @Image.operation
    def has_animation(self):
        # Animation is not supported by OpenCV
        return False


class OpenCVColorImage(BaseOpenCVImage):
    @classmethod
    def check(cls):
        super().check()
        _numpy()

    @classmethod
    @Image.converter_from(RGBImageBuffer)
    def from_buffer_rgb(cls, image_buffer):
        """
        Converts a Color Image buffer into a numpy array suitable for use with OpenCV
        """
        numpy = _numpy()
        cv2 = _cv2()

        image = numpy.frombuffer(image_buffer.data, dtype=numpy.uint8)
        image = image.reshape(image_buffer.size[1], image_buffer.size[0], 3)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        return cls(image, image_buffer.size)


class OpenCVGrayscaleImage(BaseOpenCVImage):
    face_haar_flags = 0
    face_min_neighbors = 3
    face_haar_scale = 1.1
    face_min_size = (40, 40)

    @Image.operation
    def detect_features(self):
        """
        Find interesting features of an image suitable for cropping to.
        """
        numpy = _numpy()
        cv2 = _cv2()
        points = cv2.goodFeaturesToTrack(self.image, 20, 0.04, 1.0)
        if points is None:
            return []
        else:
            points = numpy.reshape(
                points, (-1, 2)
            )  # Numpy returns it with an extra third dimension
            return points.tolist()

    @Image.operation
    def detect_faces(self, cascade_filename="haarcascade_frontalface_alt2.xml"):
        """
        Run OpenCV face detection on the image. Returns a list of coordinates representing a box around each face.
        """
        cv2 = _cv2()
        cascade_filename = self._find_cascade(cascade_filename)
        cascade = cv2.CascadeClassifier(cascade_filename)
        equalised_image = cv2.equalizeHist(self.image)
        faces = cascade.detectMultiScale(
            equalised_image,
            self.face_haar_scale,
            self.face_min_neighbors,
            self.face_haar_flags,
            self.face_min_size,
        )
        return [
            (
                face[0],
                face[1],
                face[0] + face[2],
                face[1] + face[3],
            )
            for face in faces
        ]

    def _find_cascade(self, cascade_filename):
        """
        Find the requested OpenCV cascade file.  If a relative path was provided, check local cascades directory.
        """
        if not os.path.isabs(cascade_filename):
            cascade_filename = os.path.join(
                os.path.dirname(os.path.dirname(__file__)),
                "data/cascades",
                cascade_filename,
            )
        return cascade_filename

    @classmethod
    @Image.converter_from(OpenCVColorImage)
    def from_color(cls, colour_image):
        """
        Convert OpenCVColorImage to an OpenCVGrayscaleImage.
        """
        cv2 = _cv2()
        image = cv2.cvtColor(colour_image.image, cv2.COLOR_BGR2GRAY)
        return cls(image, colour_image.size)


willow_image_classes = [OpenCVColorImage, OpenCVGrayscaleImage]