khushalps
a7c5c98bb8
git-subtree-dir: utils_v2 git-subtree-split: 13ad1588815ba34a57089927a5618f683d44cd29
224 lines
9.1 KiB
Python
224 lines
9.1 KiB
Python
"""
|
|
|
|
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())
|