|
import argparse |
|
import os |
|
from typing import Optional, Callable |
|
|
|
import torch |
|
from PIL import Image |
|
from torch import Tensor, nn |
|
from torchvision.transforms import transforms, functional |
|
from tqdm.auto import tqdm |
|
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation |
|
|
|
DEVICE = "cuda" |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser(description="ClipSeg script.") |
|
parser.add_argument( |
|
"--sample_dir", |
|
type=str, |
|
required=True, |
|
help="directory where samples are located", |
|
) |
|
parser.add_argument( |
|
"--add_prompt", |
|
type=str, |
|
required=True, |
|
action="append", |
|
help="a prompt used to create a mask", |
|
dest="prompts", |
|
) |
|
parser.add_argument( |
|
"--mode", |
|
type=str, |
|
default='fill', |
|
required=False, |
|
help="Either replace, fill, add or subtract", |
|
) |
|
parser.add_argument( |
|
"--threshold", |
|
type=float, |
|
default='0.3', |
|
required=False, |
|
help="threshold for including pixels in the mask", |
|
) |
|
parser.add_argument( |
|
"--smooth_pixels", |
|
type=int, |
|
default=5, |
|
required=False, |
|
help="radius of a smoothing operation applied to the generated mask", |
|
) |
|
parser.add_argument( |
|
"--expand_pixels", |
|
type=int, |
|
default=10, |
|
required=False, |
|
help="amount of expansion of the generated mask in all directions", |
|
) |
|
|
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
class MaskSample: |
|
def __init__(self, filename: str): |
|
self.image_filename = filename |
|
self.mask_filename = os.path.splitext(filename)[0] + "-masklabel.png" |
|
|
|
self.image = None |
|
self.mask_tensor = None |
|
|
|
self.height = 0 |
|
self.width = 0 |
|
|
|
self.image2Tensor = transforms.Compose([ |
|
transforms.ToTensor(), |
|
]) |
|
|
|
self.tensor2Image = transforms.Compose([ |
|
transforms.ToPILImage(), |
|
]) |
|
|
|
def get_image(self) -> Image: |
|
if self.image is None: |
|
self.image = Image.open(self.image_filename).convert('RGB') |
|
self.height = self.image.height |
|
self.width = self.image.width |
|
|
|
return self.image |
|
|
|
def get_mask_tensor(self) -> Tensor: |
|
if self.mask_tensor is None and os.path.exists(self.mask_filename): |
|
mask = Image.open(self.mask_filename).convert('L') |
|
mask = self.image2Tensor(mask) |
|
mask = mask.to(DEVICE) |
|
self.mask_tensor = mask.unsqueeze(0) |
|
|
|
return self.mask_tensor |
|
|
|
def set_mask_tensor(self, mask_tensor: Tensor): |
|
self.mask_tensor = mask_tensor |
|
|
|
def add_mask_tensor(self, mask_tensor: Tensor): |
|
mask = self.get_mask_tensor() |
|
if mask is None: |
|
mask = mask_tensor |
|
else: |
|
mask += mask_tensor |
|
mask = torch.clamp(mask, 0, 1) |
|
|
|
self.mask_tensor = mask |
|
|
|
def subtract_mask_tensor(self, mask_tensor: Tensor): |
|
mask = self.get_mask_tensor() |
|
if mask is None: |
|
mask = mask_tensor |
|
else: |
|
mask -= mask_tensor |
|
mask = torch.clamp(mask, 0, 1) |
|
|
|
self.mask_tensor = mask |
|
|
|
def save_mask(self): |
|
if self.mask_tensor is not None: |
|
mask = self.mask_tensor.cpu().squeeze() |
|
mask = self.tensor2Image(mask).convert('RGB') |
|
mask.save(self.mask_filename) |
|
|
|
|
|
class ClipSeg: |
|
def __init__(self): |
|
self.processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") |
|
|
|
self.model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") |
|
self.model.eval() |
|
self.model.to(DEVICE) |
|
|
|
self.smoothing_kernel_radius = None |
|
self.smoothing_kernel = self.__create_average_kernel(self.smoothing_kernel_radius) |
|
|
|
self.expand_kernel_radius = None |
|
self.expand_kernel = self.__create_average_kernel(self.expand_kernel_radius) |
|
|
|
@staticmethod |
|
def __create_average_kernel(kernel_radius: Optional[int]): |
|
if kernel_radius is None: |
|
return None |
|
|
|
kernel_size = kernel_radius * 2 + 1 |
|
kernel_weights = torch.ones(1, 1, kernel_size, kernel_size) / (kernel_size * kernel_size) |
|
kernel = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=kernel_size, bias=False, padding_mode='replicate', padding=kernel_radius) |
|
kernel.weight.data = kernel_weights |
|
kernel.requires_grad_(False) |
|
kernel.to(DEVICE) |
|
return kernel |
|
|
|
@staticmethod |
|
def __get_sample_filenames(sample_dir: str) -> [str]: |
|
filenames = [] |
|
for filename in os.listdir(sample_dir): |
|
ext = os.path.splitext(filename)[1].lower() |
|
if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp'] and '-masklabel.png' not in filename: |
|
filenames.append(os.path.join(sample_dir, filename)) |
|
|
|
return filenames |
|
|
|
def __process_mask(self, mask: Tensor, target_height: int, target_width: int, threshold: float) -> Tensor: |
|
while len(mask.shape) < 4: |
|
mask = mask.unsqueeze(0) |
|
|
|
mask = torch.sigmoid(mask) |
|
mask = mask.sum(1).unsqueeze(1) |
|
if self.smoothing_kernel is not None: |
|
mask = self.smoothing_kernel(mask) |
|
mask = functional.resize(mask, [target_height, target_width]) |
|
mask = (mask > threshold).float() |
|
if self.expand_kernel is not None: |
|
mask = self.expand_kernel(mask) |
|
mask = (mask > 0).float() |
|
|
|
return mask |
|
|
|
def mask_image(self, filename: str, prompts: [str], mode: str = 'fill', threshold: float = 0.3, smooth_pixels: int = 5, expand_pixels: int = 10): |
|
""" |
|
Masks a sample |
|
|
|
Parameters: |
|
filename (`str`): a sample filename |
|
prompts (`[str]`): a list of prompts used to create a mask |
|
mode (`str`): can be one of |
|
- replace: creates new masks for all samples, even if a mask already exists |
|
- fill: creates new masks for all samples without a mask |
|
- add: adds the new region to existing masks |
|
- subtract: subtracts the new region from existing masks |
|
threshold (`float`): threshold for including pixels in the mask |
|
smooth_pixels (`int`): radius of a smoothing operation applied to the generated mask |
|
expand_pixels (`int`): amount of expansion of the generated mask in all directions |
|
""" |
|
|
|
mask_sample = MaskSample(filename) |
|
|
|
if mode == 'fill' and mask_sample.get_mask_tensor() is not None: |
|
return |
|
|
|
if self.smoothing_kernel_radius != smooth_pixels: |
|
self.smoothing_kernel = self.__create_average_kernel(smooth_pixels) |
|
self.smoothing_kernel_radius = smooth_pixels |
|
|
|
if self.expand_kernel_radius != expand_pixels: |
|
self.expand_kernel = self.__create_average_kernel(expand_pixels) |
|
self.expand_kernel_radius = expand_pixels |
|
|
|
inputs = self.processor(text=prompts, images=[mask_sample.get_image()] * len(prompts), padding="max_length", return_tensors="pt") |
|
inputs.to(DEVICE) |
|
with torch.no_grad(): |
|
outputs = self.model(**inputs) |
|
predicted_mask = self.__process_mask(outputs.logits, mask_sample.height, mask_sample.width, threshold) |
|
|
|
if mode == 'replace' or mode == 'fill': |
|
mask_sample.set_mask_tensor(predicted_mask) |
|
elif mode == 'add': |
|
mask_sample.add_mask_tensor(predicted_mask) |
|
elif mode == 'subtract': |
|
mask_sample.subtract_mask_tensor(predicted_mask) |
|
|
|
mask_sample.save_mask() |
|
|
|
def mask_folder( |
|
self, |
|
sample_dir: str, |
|
prompts: [str], |
|
mode: str = 'fill', |
|
threshold: float = 0.3, |
|
smooth_pixels: int = 5, |
|
expand_pixels: int = 10, |
|
progress_callback: Callable[[int, int], None] = None, |
|
error_callback: Callable[[str], None] = None, |
|
): |
|
""" |
|
Masks all samples in a folder |
|
|
|
Parameters: |
|
sample_dir (`str`): directory where samples are located |
|
prompts (`[str]`): a list of prompts used to create a mask |
|
mode (`str`): can be one of |
|
- replace: creates new masks for all samples, even if a mask already exists |
|
- fill: creates new masks for all samples without a mask |
|
- add: adds the new region to existing masks |
|
- subtract: subtracts the new region from existing masks |
|
threshold (`float`): threshold for including pixels in the mask |
|
smooth_pixels (`int`): radius of a smoothing operation applied to the generated mask |
|
expand_pixels (`int`): amount of expansion of the generated mask in all directions |
|
progress_callback (`Callable[[int, int], None]`): called after every processed image |
|
error_callback (`Callable[[str], None]`): called for every exception |
|
""" |
|
|
|
filenames = self.__get_sample_filenames(sample_dir) |
|
self.mask_images( |
|
filenames=filenames, |
|
prompts=prompts, |
|
mode=mode, |
|
threshold=threshold, |
|
smooth_pixels=smooth_pixels, |
|
expand_pixels=expand_pixels, |
|
progress_callback=progress_callback, |
|
error_callback=error_callback, |
|
) |
|
|
|
def mask_images( |
|
self, |
|
filenames: [str], |
|
prompts: [str], |
|
mode: str = 'fill', |
|
threshold: float = 0.3, |
|
smooth_pixels: int = 5, |
|
expand_pixels: int = 10, |
|
progress_callback: Callable[[int, int], None] = None, |
|
error_callback: Callable[[str], None] = None, |
|
): |
|
""" |
|
Masks all samples in a list |
|
|
|
Parameters: |
|
filenames (`[str]`): a list of sample filenames |
|
prompts (`[str]`): a list of prompts used to create a mask |
|
mode (`str`): can be one of |
|
- replace: creates new masks for all samples, even if a mask already exists |
|
- fill: creates new masks for all samples without a mask |
|
- add: adds the new region to existing masks |
|
- subtract: subtracts the new region from existing masks |
|
threshold (`float`): threshold for including pixels in the mask |
|
smooth_pixels (`int`): radius of a smoothing operation applied to the generated mask |
|
expand_pixels (`int`): amount of expansion of the generated mask in all directions |
|
progress_callback (`Callable[[int, int], None]`): called after every processed image |
|
error_callback (`Callable[[str], None]`): called for every exception |
|
""" |
|
|
|
if progress_callback is not None: |
|
progress_callback(0, len(filenames)) |
|
for i, filename in enumerate(tqdm(filenames)): |
|
try: |
|
self.mask_image(filename, prompts, mode, threshold, smooth_pixels, expand_pixels) |
|
except Exception as e: |
|
if error_callback is not None: |
|
error_callback(filename) |
|
if progress_callback is not None: |
|
progress_callback(i + 1, len(filenames)) |
|
|
|
|
|
def main(): |
|
args = parse_args() |
|
clip_seg = ClipSeg() |
|
clip_seg.mask_folder( |
|
sample_dir=args.sample_dir, |
|
prompts=args.prompts, |
|
mode=args.mode, |
|
threshold=args.threshold, |
|
smooth_pixels=args.smooth_pixels, |
|
expand_pixels=args.expand_pixels, |
|
error_callback=lambda filename: print("Error while processing image " + filename) |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
