Squashed 'utils_v2/' content from commit f03179d
git-subtree-dir: utils_v2 git-subtree-split: f03179d339e69fdcdff2f0cc06e1f352024c55d3
This commit is contained in:
yatmesh
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
|
||||
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())
|
||||
Reference in New Issue
Block a user