"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Saturday, 14th Sept., 2024

    OBJECTIVE:

        To provide an easy way to get masks from dichotomous image segmentation.

        This code uses a very specific model from HuggingFace: "ZhengPeng7/BiRefNet-portrait".
        You may experiment with other behaviour_models too, but make sure that model is made for "dichotomous" behaviour. This
        means that the model should have only two classes like "foreground", and "background". The specified model was
        trained for implementing portrait mode style blurring of backgrounds.

        The originally tested model has an MIT license as per their GitHub page. The code in this file may or may not
        support drop-in replacement for other behaviour_models, please be aware about this.

    REFERENCES:

        01. https://huggingface.co/ZhengPeng7/BiRefNet-portrait
        02. https://github.com/ZhengPeng7/BiRefNet

    DOWNLOADS:

        01. https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# For system-level activities:
import io

# To work with PIL images:
from PIL import Image

# To make asynchronous HTTP calls:
from shared.variables import http_client

# For using the AI model:
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  CLASSES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class DichotomousSegmenter:

    def __init__(
            self,
            model = r"zhengpeng7/BiRefNet-portrait"
    ):

        """
        This class uses Dichotomous Image Segmentation to produce a mask of what is in the foreground (or what is of
        interest in a given image).
        :param model: The path of the model.
        """

        # Check if GPU is available for faster predictions:
        self.__cuda_is_available = True if torch.cuda.is_available() else False

        # Initialize the model:
        self.__model = AutoModelForImageSegmentation.from_pretrained(
            pretrained_model_name_or_path = model,
            trust_remote_code = True
        )
        torch.set_float32_matmul_precision(["high", "highest"][0])
        self.__model.to("cuda" if self.__cuda_is_available else "cpu")

    @staticmethod
    async def read_image(source):

        """
        Reads an image in whichever format it is provided and returns it as a PIL object.
        :param source: The image as either a path or a URL or a io.BytesIO object.
        :return: The image opened as a PIL object.
        """

        # If the input image is already a PIL image:
        if isinstance(source, Image.Image): return source

        # If a buffer is provided:
        elif isinstance(source, io.BytesIO):
            source.seek(0)
            return Image.open(source)

        # If a string is provided (local path or URL):
        elif isinstance(source, str):
            if source.startswith("http://") or source.startswith("https://"):
                response = await http_client.get(source)
                return Image.open(io.BytesIO(response.content))
            else: return Image.open(source)

    async def get_mask(self, image):

        """
        Process the image to figure out the mask. In some cases (like the image of a sunset) you will have no white
        pixels in the mask. White pixels represent the areas that the AI considered to be the foreground (area of
        interest) and the black parts are the background.
        :param image: The image data as either a PIL object, or a path to a file on the local disk, or a URL.
        :return: The mask as a PIL object.
        """

        # Prepare the image transformer:
        image_size = (1024, 1024)
        transform_image = transforms.Compose([
            transforms.Resize(image_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        # Read and transform the image:
        image = await self.read_image(image)
        input_image = transform_image(image).unsqueeze(0).to("cuda" if self.__cuda_is_available else "cpu")

        # Make the prediction:
        with torch.no_grad():
            if self.__cuda_is_available: mask = self.__model(input_image)[-1].sigmoid().cuda()
            else: mask = self.__model(input_image)[-1].sigmoid().cpu()

        # Create a PIL object of the mask data:
        mask = mask[0].squeeze()
        mask = transforms.ToPILImage()(mask)
        mask = mask.resize(image.size)

        # Done here:
        return mask


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    import asyncio
    import numpy as np

    async def main():

        model = DichotomousSegmenter()

        images = [
            # r"/home/developer/Downloads/kate.jpg",
            # r"/home/developer/Downloads/IMG-20240730-WA0001.jpg",
            # r"/home/developer/Downloads/pexels-pixabay-57416.jpg",
            # r"/home/developer/Downloads/summer_clothes.jpg",
            # r"/home/developer/Downloads/low_res_cat - upscaled.png",
            # r"/home/developer/Downloads/blurry_traffic.jpg",
            r"/home/developer/Downloads/card_0.jpg",
            r"/home/developer/Downloads/card_1.jpg",
            r"/home/developer/Downloads/card_2.jpg",
            r"/home/developer/Downloads/card_3.jpg",
            r"/home/developer/Downloads/sushmita_card.jpg",
            r"/home/developer/Downloads/niranjan_card.jpg",
            r"/home/developer/Downloads/niranjan_card_2.jpg",
        ]

        for image in images:
            print(image.split("/")[-1])
            image = await DichotomousSegmenter.read_image(image)
            mask = await model.get_mask(image)
            image.show()
            mask.show()
            print("MASK!")
            array = np.array(mask)
            print("Dimensions of the array:", array.shape)
            masked_count = np.sum(array > 125)
            print(masked_count, type(masked_count), np.size, type(np.size))
            percentage = masked_count / np.prod(array.shape)
            print("Area:", round(percentage * 100, 4))
            # image.putalpha(mask)
            # image.show()
            # mask.show()


    asyncio.run(main())
