faceSym.py

# based on eggplants/face-symmetrizer
from __future__ import annotations

import io
import re
from copy import copy
from os import path
from typing import Any, Dict, List, Tuple
from urllib.request import urlopen

import face_recognition  # type: ignore[import]
import numpy as np
from PIL import Image, ImageDraw, ImageOps

PILImage = Image.Image
FaceLandmarks = List[Dict[str, List[Tuple[Any, ...]]]]


class FaceIsNotDetected(Exception):
    """[summary]

    Args:
        Exception ([type]): [description]
    """

    pass


class FaceSym:
    """[summary]"""

    SimImages = Tuple[PILImage, PILImage, PILImage, PILImage, PILImage, PILImage]

    def __init__(self, img_location: str) -> None:
        """[summary]

        Args:
            img_location (str): [description]

        Raises:
            ValueError: [description]
        """
        self.f_img: np.ndarray[Any, Any]
        self.image_location = img_location
        if self.__is_valid_url(img_location):
            self.__load_from_url(img_location)
        elif path.isfile(img_location):
            self.__load_from_local(img_location)
        else:
            raise ValueError(
                f"{repr(img_location)} is not a valid location of an image."
            )

        self.f_img_PIL = Image.fromarray(self.f_img)
        self.image_size: tuple[int, int] = self.f_img_PIL.size
        self.face_locations = face_recognition.face_locations(self.f_img)
        self.face_landmarks = face_recognition.face_landmarks(self.f_img)
        self.mid_face_locations = self.__get_mid_face_locations(self.face_landmarks)
        self.face_count = len(self.face_locations)

    def get_cropped_face_images(self,) -> list[PILImage]:
        """[summary]

        Returns:
            List[PILImage]: [description]
        """
        images = []
        for face_location in self.face_locations:
            top, right, bottom, left = face_location
            cropped_face_img = self.f_img[top:bottom, left:right]
            pil_img = Image.fromarray(cropped_face_img)
            
            images.append(pil_img)

        return images

    def get_face_box_drawed_image(self) -> PILImage:
        """[summary]

        Returns:
            PILImage: [description]
        """
        pil = copy(self.f_img_PIL)
        draw = ImageDraw.Draw(pil)
        for idx, (top, right, bottom, left) in enumerate(self.face_locations):
            name = str(f"{idx:02d}")
            mid_face = self.mid_face_locations[idx]

            draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))

            _, text_height = draw.textsize(name)
            draw.rectangle(
                ((left, bottom - text_height - 10), (right, bottom)),
                fill=(0, 0, 255),
                outline=(0, 0, 255),
            )
            draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255))

            draw.line(
                ((mid_face[0], -10), mid_face, (mid_face[0], self.image_size[0])),
                fill=(255, 255, 0),
                width=10,
            )
        del draw
        return pil

    def get_full_image(
        self, is_pil: bool = False
    ) -> np.ndarray[Any, Any] | PILImage:
        """[summary]

        Args:
            is_pil (bool, optional): [description]. Defaults to False.

        Returns:
            Union[np.ndarray, PILImage]: [description]
        """
        
        if is_pil:
            return self.f_img_PIL
        else:
            return self.f_img

    def get_symmetrized_images(self, idx: int = 0) -> SimImages:
        """[summary]

        Args:
            idx (int, optional): [description]. Defaults to 0.

        Returns:
            SimImages: [description]
        """

        def get_concat_h(im1: PILImage, im2: PILImage) -> PILImage:
            dst = Image.new("RGB", (im1.width + im2.width, im1.height))
            dst.paste(im1, (0, 0))
            dst.paste(im2, (im1.width, 0))
            return dst

        face_count = len(self.mid_face_locations)
        if face_count < 1:
            raise FaceIsNotDetected
        elif face_count <= idx:
            raise IndexError(f"0 <= idx <= {face_count - 1}")
        else:
            mid_face = self.mid_face_locations[idx]

        cropped_left_img = self.f_img[0 : self.image_size[1], 0 : int(mid_face[0])]
        cropped_right_img = self.f_img[
            0 : self.image_size[1], int(mid_face[0]) : self.image_size[0]
        ]

        pil_img_left = Image.fromarray(cropped_left_img)
        pil_img_left_mirrored = ImageOps.mirror(pil_img_left)
        pil_img_left_inner = get_concat_h(pil_img_left, pil_img_left_mirrored)
        pil_img_left_outer = get_concat_h(pil_img_left_mirrored, pil_img_left)

        pil_img_right = Image.fromarray(cropped_right_img)
        pil_img_right_mirrored = ImageOps.mirror(pil_img_right)
        pil_img_right_inner = get_concat_h(pil_img_right_mirrored, pil_img_right)
        pil_img_right_outer = get_concat_h(pil_img_right, pil_img_right_mirrored)

        return (
            pil_img_left,
            pil_img_left_inner,
            pil_img_left_outer,
            pil_img_right,
            pil_img_right_inner,
            pil_img_right_outer,
        )

    def __load_from_url(self, url: str) -> None:
        """[summary]

        Args:
            url (str): [description]

        Raises:
            ValueError: [description]
        """
        if not self.__is_valid_url(url):
            raise ValueError(f"{repr(url)} is not valid url")
        else:
            img_data = io.BytesIO(urlopen(url).read())
            self.f_img = face_recognition.load_image_file(img_data)

    def __load_from_local(self, path_: str) -> None:
        if path.isfile(path_):
            self.f_img = face_recognition.load_image_file(path_)

    @staticmethod
    def __is_valid_url(url: str) -> bool:
        """[summary]

        Args:
            url (str): [description]

        Returns:
            bool: [description]

        Note:
            Copyright (c) Django Software Foundation and individual
            contributors. All rights reserved.
        """
        regex = re.compile(
            r"^(?:http|ftp)s?://"  # http:// or https://
            # domain...
            r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|"
            r"[A-Z0-9-]{2,}\.?)|"
            r"localhost|"  # localhost...
            r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})"  # ...or ip
            r"(?::\d+)?"  # optional port
            r"(?:/?|[/?]\S+)$",
            re.IGNORECASE,
        )
        return re.match(regex, url) is not None

    @staticmethod
    def __get_mid_face_locations(
        face_landmarks: FaceLandmarks,
    ) -> list[tuple[int, int]]:
        """[summary]

        Args:
            face_landmarks (FaceLandmarks): [description]

        Returns:
            List[Tuple[int, int]]: [description]
        """

        def mean(lst: list[int]) -> int:
            return int(sum(lst) / len(lst))

        mid_faces = []
        for face_landmark in face_landmarks:
            if not ("left_eye" in face_landmark and "right_eye" in face_landmark):
                raise ValueError("eye locations was missing.")
            l_e_xs = [i[0] for i in face_landmark["left_eye"]]
            l_e_ys = [i[1] for i in face_landmark["left_eye"]]
            r_e_xs = [i[0] for i in face_landmark["right_eye"]]
            r_e_ys = [i[1] for i in face_landmark["right_eye"]]
            mid_face = (
                (mean(l_e_xs) + mean(r_e_xs)) // 2,
                (mean(l_e_ys) + mean(r_e_ys)) // 2,
            )
            mid_faces.append(mid_face)
        return mid_faces


def main() -> None:
    """[summary]"""
    data = list(
        map(
            lambda x: "https://pbs.twimg.com/media/%s?format=jpg" % x,
            [
                "E7okHDEVUAE1O6i",
                "E7jaibgUcAUWvg-",
                "E7jahEbUcAMNLdU",
                "E7Jqli9VEAEStvs",
                "E7Jqk-aUcAcfg3o",
                "E7EhGi2XoAsMrO5",
                "E5dhLccUYAUD5Yx",
                "E5TOAqUVUAMckXT",
                "E4vK6e0VgAAksnK",
                "E4Va7u4VkAAKde3",
                "E4A0ksEUYAIpynP",
                "E3xXzcyUYAIX1dC",
                "E2zkvONVcAQEE_S",
                "E1cBsxDUcAIe_LZ",
                "E1W4HTRVUAgYkmo",
                "E1HbVAeVIAId5yP",
                "E09INVFUcAYpcWo",
                "E0oh0hmUUAAfJV9",
            ],
        )
    )
    success, fail = 0, 0
    for idx, link in enumerate(data):
        print(f"[{idx:02d}]", link, end="")
        f = FaceSym(link)
        if f.face_count != 0:
            print("=>Detected")
            f.get_symmetrized_images()
            success += 1
        else:
            print("=>Not Detected")
            fail += 1

    else:
        print(f"DATA: {len(data)}", f"OK: {success}", f"NG: {fail}")


if __name__ == "__main__":
    main()