diff --git a/gfrbmg2plus.py b/gfrbmg2plus.py new file mode 100644 index 0000000..a4eacfa --- /dev/null +++ b/gfrbmg2plus.py @@ -0,0 +1,213 @@ +import os +import torch +import folder_paths +from PIL import Image +from transformers import AutoModelForImageSegmentation +from torchvision.transforms.functional import normalize +import numpy as np +import cv2 +import torch.nn.functional as F +from huggingface_hub import hf_hub_download, HfApi + +device = "cuda" if torch.cuda.is_available() else "cpu" + +class GFrbmg2Plus: + def __init__(self): + self.model = None + self.model_path = os.path.join(folder_paths.models_dir, "RMBG", "RMBG-2.0") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "invert_mask": ("BOOLEAN", {"default": False}), + "chroma_key_color": (["Black", "White", "Green", "Red", "Blue", "Gray"], {"default": "Black"}), + "postprocess_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 20.0, "step": 0.1}), + "edge_enhancement": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 50.0, "step": 0.1}), + "blur_edges": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 50.0, "step": 0.5}), + "expand_mask": ("FLOAT", {"default": 0.0, "min": -50.0, "max": 50.0, "step": 0.1}), + "normalize_image": ("BOOLEAN", {"default": True}), + "range_low": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05}), + "range_high": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}), + "contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}) + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK") + RETURN_NAMES = ("image_rgba", "image", "mask_rgb", "mask") + FUNCTION = "remove_background" + CATEGORY = "🐡 GorillaFrame/Image" + + def initialize_model(self): + if not os.path.exists(self.model_path): + os.makedirs(self.model_path, exist_ok=True) + + # ΠŸΠΎΠ»ΡƒΡ‡Π°Π΅ΠΌ список всСх Ρ„Π°ΠΉΠ»ΠΎΠ² Π² Ρ€Π΅ΠΏΠΎΠ·ΠΈΡ‚ΠΎΡ€ΠΈΠΈ + api = HfApi() + files = api.list_repo_files("briaai/RMBG-2.0") + + # Π‘ΠΊΠ°Ρ‡ΠΈΠ²Π°Π΅ΠΌ Ρ‚ΠΎΠ»ΡŒΠΊΠΎ Ρ„Π°ΠΉΠ»Ρ‹ ΠΈΠ· ΠΊΠΎΡ€Π½Π΅Π²ΠΎΠΉ Π΄ΠΈΡ€Π΅ΠΊΡ‚ΠΎΡ€ΠΈΠΈ + for file in files: + # ΠŸΡ€ΠΎΠ²Π΅Ρ€ΡΠ΅ΠΌ Ρ‡Ρ‚ΠΎ Ρ„Π°ΠΉΠ» Π² ΠΊΠΎΡ€Π½Π΅ (Π½Π΅Ρ‚ слСшСй Π² ΠΏΡƒΡ‚ΠΈ) + if '/' not in file: + if file.endswith(('.json', '.py', '.safetensors')): + print(f"Downloading {file}...") + hf_hub_download( + repo_id="briaai/RMBG-2.0", + filename=file, + local_dir=self.model_path + ) + + self.model = AutoModelForImageSegmentation.from_pretrained( + self.model_path, + trust_remote_code=True, + local_files_only=True + ) + self.model.to(device) + self.model.eval() + + def tensor2pil(self, image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + + def pil2tensor(self, image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + def resize_image(self, image): + image = image.convert('RGB') + return image.resize((1024, 1024), Image.BILINEAR) + + def normalize_mask(self, mask, range_low, range_high): + mask_np = np.array(mask, dtype=np.float32) / 255.0 + mask_np = (mask_np - range_low) / (range_high - range_low) + mask_np = np.clip(mask_np, 0, 1) + return Image.fromarray((mask_np * 255).astype(np.uint8)) + + def adjust_contrast(self, mask, contrast): + if contrast == 1.0: + return mask + mask_np = np.array(mask, dtype=np.float32) / 255.0 + mask_np = np.clip((mask_np - 0.5) * contrast + 0.5, 0, 1) + return Image.fromarray((mask_np * 255).astype(np.uint8)) + + def clean_mask(self, mask, strength, edge_enhancement, blur_edges, expand_mask): + if strength == 0 and edge_enhancement == 0 and blur_edges == 0 and expand_mask == 0: + return mask + + mask_np = np.array(mask) + + if expand_mask != 0: + kernel_size = int(5 * abs(expand_mask)) + kernel_size = max(1, kernel_size) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = np.ones((kernel_size, kernel_size), np.uint8) + if expand_mask > 0: + mask_np = cv2.dilate(mask_np, kernel, iterations=1) + else: + mask_np = cv2.erode(mask_np, kernel, iterations=1) + + if edge_enhancement > 0: + lower_threshold = int(100 * edge_enhancement) + upper_threshold = int(200 * edge_enhancement) + edges = cv2.Canny(mask_np, lower_threshold, upper_threshold) + if blur_edges > 0: + blur_size = int(5 * blur_edges) + if blur_size % 2 == 0: + blur_size += 1 + edges = cv2.GaussianBlur(edges, (blur_size, blur_size), 0) + mask_np = cv2.bitwise_or(mask_np, edges) + + if strength > 0: + kernel_size = int(5 * strength) + kernel_size = max(1, kernel_size) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = np.ones((kernel_size, kernel_size), np.uint8) + iterations = int(strength) + original_mask = mask_np.copy() + mask_smoothed = cv2.medianBlur(mask_np, kernel_size) + mask_dilated = cv2.dilate(mask_smoothed, kernel, iterations=iterations) + mask_eroded = cv2.erode(mask_dilated, kernel, iterations=iterations) + mask_np = cv2.bitwise_and(mask_eroded, original_mask) + + return Image.fromarray(mask_np) + + def remove_background(self, image, invert_mask, chroma_key_color, postprocess_strength, edge_enhancement, blur_edges, expand_mask, normalize_image, range_low, range_high, contrast): + if self.model is None: + self.initialize_model() + + processed_images = [] + processed_blacks = [] + processed_masks = [] + processed_masks_rgb = [] + + for img in image: + orig_image = self.tensor2pil(img) + w, h = orig_image.size + image = self.resize_image(orig_image) + + im_np = np.array(image) + im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1) + im_tensor = torch.unsqueeze(im_tensor, 0) + im_tensor = torch.divide(im_tensor, 255.0) + im_tensor = normalize(im_tensor, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + + if torch.cuda.is_available(): + im_tensor = im_tensor.cuda() + + with torch.no_grad(): + result = self.model(im_tensor)[-1].sigmoid().cpu() + + result = result[0].squeeze() + result = F.interpolate(result.unsqueeze(0).unsqueeze(0), size=(h, w), mode='bilinear', align_corners=True).squeeze() + + mask_pil = self.tensor2pil(result) + mask_pil = self.clean_mask(mask_pil, postprocess_strength, edge_enhancement, blur_edges, expand_mask) + + if normalize_image: + mask_pil = self.normalize_mask(mask_pil, range_low, range_high) + mask_pil = self.adjust_contrast(mask_pil, contrast) + + if invert_mask: + mask_np = np.array(mask_pil) + mask_np = 255 - mask_np + mask_pil = Image.fromarray(mask_np) + + rgba_image = orig_image.copy() + rgba_image.putalpha(mask_pil) + + background_color = { + "Black": (0, 0, 0), + "White": (255, 255, 255), + "Green": (0, 255, 0), + "Red": (255, 0, 0), + "Blue": (0, 0, 255), + "Gray": (128, 128, 128) + }[chroma_key_color] + + background_image = Image.new('RGB', orig_image.size, background_color) + background_image.paste(orig_image, mask=mask_pil) + + mask_rgb = Image.new('RGB', orig_image.size, (0, 0, 0)) + mask_rgb.paste(Image.fromarray(np.array(mask_pil).astype(np.uint8)), mask=mask_pil) + + processed_images.append(self.pil2tensor(rgba_image)) + processed_blacks.append(self.pil2tensor(background_image)) + processed_masks.append(self.pil2tensor(mask_pil)) + processed_masks_rgb.append(self.pil2tensor(mask_rgb)) + + return ( + torch.cat(processed_images, dim=0), + torch.cat(processed_blacks, dim=0), + torch.cat(processed_masks_rgb, dim=0), + torch.cat(processed_masks, dim=0) + ) + +NODE_CLASS_MAPPINGS = { + "GFrbmg2Plus": GFrbmg2Plus +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GFrbmg2Plus": "🐡 GF Remove Background (GFrbmg2Plus)" +} \ No newline at end of file