226 lines
9.8 KiB
Python
226 lines
9.8 KiB
Python
import os
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, List, Tuple, Optional
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
from huggingface_hub import hf_hub_download
|
|
import folder_paths
|
|
|
|
from AILab_utils import get_device, clean_vram, unload_model as release_model, tensor2pil, pil2tensor
|
|
|
|
device = get_device()
|
|
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
|
|
|
|
MODEL_SUBDIR = os.path.join(folder_paths.models_dir, "RMBG", "finegrain-box-segmenter")
|
|
REPO_ID = "1038lab/finegrain-box-segmenter"
|
|
WEIGHTS_FILENAME = "finegrain-box-segmenter.safetensors"
|
|
|
|
|
|
def parse_bounding_boxes(bboxes_input: Any, box_str: str, img_w: int, img_h: int) -> List[Tuple[int, int, int, int]]:
|
|
boxes = []
|
|
|
|
if bboxes_input is not None:
|
|
raw_list = []
|
|
if isinstance(bboxes_input, str):
|
|
cleaned = bboxes_input.strip()
|
|
if cleaned:
|
|
try:
|
|
parsed = json.loads(cleaned)
|
|
if isinstance(parsed, list):
|
|
raw_list = parsed
|
|
elif isinstance(parsed, dict):
|
|
raw_list = parsed.get("bboxes", parsed.get("boxes", []))
|
|
except Exception:
|
|
pass
|
|
elif isinstance(bboxes_input, dict):
|
|
raw_list = bboxes_input.get("bboxes", bboxes_input.get("boxes", []))
|
|
elif isinstance(bboxes_input, (list, tuple)):
|
|
if len(bboxes_input) == 4 and all(isinstance(x, (int, float)) for x in bboxes_input):
|
|
raw_list = [bboxes_input]
|
|
else:
|
|
raw_list = bboxes_input
|
|
|
|
for item in raw_list:
|
|
if isinstance(item, (list, tuple)) and len(item) >= 4:
|
|
x1, y1, x2, y2 = float(item[0]), float(item[1]), float(item[2]), float(item[3])
|
|
if max(x1, y1, x2, y2) <= 1.0:
|
|
x1, y1, x2, y2 = x1 * img_w, y1 * img_h, x2 * img_w, y2 * img_h
|
|
x1_c = max(0, min(img_w, int(round(min(x1, x2)))))
|
|
y1_c = max(0, min(img_h, int(round(min(y1, y2)))))
|
|
x2_c = max(0, min(img_w, int(round(max(x1, x2)))))
|
|
y2_c = max(0, min(img_h, int(round(max(y1, y2)))))
|
|
if x2_c > x1_c and y2_c > y1_c:
|
|
boxes.append((x1_c, y1_c, x2_c, y2_c))
|
|
|
|
if not boxes and box_str:
|
|
cleaned = box_str.strip().strip("[]()").replace(" ", "")
|
|
parts = [p for p in cleaned.split(",") if p]
|
|
if len(parts) >= 4:
|
|
try:
|
|
x1, y1, x2, y2 = float(parts[0]), float(parts[1]), float(parts[2]), float(parts[3])
|
|
if max(x1, y1, x2, y2) <= 1.0:
|
|
x1, y1, x2, y2 = x1 * img_w, y1 * img_h, x2 * img_w, y2 * img_h
|
|
x1_c = max(0, min(img_w, int(round(min(x1, x2)))))
|
|
y1_c = max(0, min(img_h, int(round(min(y1, y2)))))
|
|
x2_c = max(0, min(img_w, int(round(max(x1, x2)))))
|
|
y2_c = max(0, min(img_h, int(round(max(y1, y2)))))
|
|
if x2_c > x1_c and y2_c > y1_c:
|
|
boxes.append((x1_c, y1_c, x2_c, y2_c))
|
|
except ValueError:
|
|
pass
|
|
|
|
return boxes
|
|
|
|
|
|
class AILab_FinegrainBoxSegment:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"tooltip": "Input image batch to segment with high definition 1024px quality."}),
|
|
},
|
|
"optional": {
|
|
"bboxes": ("BBOX", {"tooltip": "Bounding box(es) from Florence-2 or YOLO detection nodes."}),
|
|
"box_coordinates": ("STRING", {"default": "", "tooltip": "Manual bounding box as 'x1, y1, x2, y2' in pixel or normalized [0..1] coordinates."}),
|
|
"margin": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 0.5, "step": 0.01, "tooltip": "Safety padding margin around bounding boxes."}),
|
|
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": "Blur mask edges (0 to disable)."}),
|
|
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": "Expand (positive) or shrink (negative) mask edges in pixels."}),
|
|
"invert_output": ("BOOLEAN", {"default": False, "tooltip": "Invert the output mask."}),
|
|
"unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM immediately after execution."}),
|
|
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Background type: Alpha (transparent) or Color (custom solid background)."}),
|
|
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": "Hex color code for background when Color mode is selected."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
|
|
FUNCTION = "segment"
|
|
CATEGORY = "🧪AILab/🧽RMBG"
|
|
|
|
def __init__(self):
|
|
self.segmenter = None
|
|
self.current_weights_path = None
|
|
self.target_device = device
|
|
|
|
def _ensure_model(self, margin=0.05):
|
|
if self.segmenter is not None:
|
|
self.segmenter.margin = margin
|
|
return self.segmenter
|
|
|
|
try:
|
|
from refiners.solutions import BoxSegmenter
|
|
except ImportError as e:
|
|
raise RuntimeError("The 'refiners' library is required for Finegrain Box Segmenter. Please run: pip install refiners") from e
|
|
|
|
os.makedirs(MODEL_SUBDIR, exist_ok=True)
|
|
local_weights = os.path.join(MODEL_SUBDIR, WEIGHTS_FILENAME)
|
|
|
|
if not os.path.exists(local_weights):
|
|
print(f"[Finegrain] Downloading {WEIGHTS_FILENAME} from {REPO_ID} to {MODEL_SUBDIR}...")
|
|
local_weights = hf_hub_download(
|
|
repo_id=REPO_ID,
|
|
filename=WEIGHTS_FILENAME,
|
|
local_dir=MODEL_SUBDIR,
|
|
local_dir_use_symlinks=False
|
|
)
|
|
|
|
self.segmenter = BoxSegmenter(
|
|
margin=margin,
|
|
weights=local_weights,
|
|
device=str(self.target_device)
|
|
)
|
|
self.current_weights_path = local_weights
|
|
return self.segmenter
|
|
|
|
def segment(self, image, bboxes=None, box_coordinates="", margin=0.05,
|
|
mask_blur=0, mask_offset=0, invert_output=False, unload_model=False,
|
|
background="Alpha", background_color="#222222", **kwargs):
|
|
do_unload = kwargs.get("unload_model", unload_model)
|
|
segmenter = self._ensure_model(margin=margin)
|
|
|
|
processed_images = []
|
|
processed_masks = []
|
|
|
|
try:
|
|
for single_img in image:
|
|
pil_img = tensor2pil(single_img).convert("RGB")
|
|
orig_w, orig_h = pil_img.size
|
|
|
|
parsed_boxes = parse_bounding_boxes(bboxes, box_coordinates, orig_w, orig_h)
|
|
|
|
if not parsed_boxes:
|
|
# Segment whole image / default subject
|
|
mask_pil = segmenter(pil_img, box_prompt=None)
|
|
combined_mask_np = np.array(mask_pil, dtype=np.float32)
|
|
else:
|
|
combined_mask_np = np.zeros((orig_h, orig_w), dtype=np.float32)
|
|
for box in parsed_boxes:
|
|
mask_box = segmenter(pil_img, box_prompt=box)
|
|
mask_np = np.array(mask_box, dtype=np.float32)
|
|
combined_mask_np = np.maximum(combined_mask_np, mask_np)
|
|
|
|
if combined_mask_np.max() > 1.0:
|
|
combined_mask_np = combined_mask_np / 255.0
|
|
|
|
mask_img = Image.fromarray((np.clip(combined_mask_np, 0.0, 1.0) * 255.0).astype(np.uint8), mode="L")
|
|
|
|
if mask_offset != 0:
|
|
offset_k = abs(mask_offset) * 2 + 1
|
|
filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter
|
|
mask_img = mask_img.filter(filter_type(offset_k))
|
|
|
|
if mask_blur > 0:
|
|
mask_img = mask_img.filter(ImageFilter.GaussianBlur(radius=mask_blur))
|
|
|
|
if invert_output:
|
|
mask_inv = 255.0 - np.array(mask_img, dtype=np.float32)
|
|
mask_img = Image.fromarray(np.clip(mask_inv, 0, 255).astype(np.uint8), mode="L")
|
|
|
|
mask_tensor = pil2tensor(mask_img)
|
|
processed_masks.append(mask_tensor)
|
|
|
|
if background == "Alpha":
|
|
cutout = Image.new("RGBA", (orig_w, orig_h), (0, 0, 0, 0))
|
|
cutout.paste(pil_img, (0, 0), mask_img)
|
|
processed_images.append(pil2tensor(cutout))
|
|
else:
|
|
bg_hex = background_color.lstrip("#")
|
|
if len(bg_hex) == 6:
|
|
bg_rgb = tuple(int(bg_hex[i:i+2], 16) for i in (0, 2, 4))
|
|
else:
|
|
bg_rgb = (34, 34, 34)
|
|
|
|
bg_canvas = Image.new("RGB", (orig_w, orig_h), bg_rgb)
|
|
bg_canvas.paste(pil_img, (0, 0), mask_img)
|
|
processed_images.append(pil2tensor(bg_canvas))
|
|
|
|
out_image = torch.cat(processed_images, dim=0)
|
|
out_mask = torch.cat(processed_masks, dim=0)
|
|
out_mask_image = out_mask.unsqueeze(-1).expand(-1, -1, -1, 3)
|
|
|
|
return (out_image, out_mask, out_mask_image)
|
|
|
|
finally:
|
|
if do_unload:
|
|
if self.segmenter is not None:
|
|
if hasattr(self.segmenter, "model") and hasattr(self.segmenter.model, "to"):
|
|
try:
|
|
self.segmenter.model.to("cpu")
|
|
except Exception:
|
|
pass
|
|
self.segmenter = None
|
|
clean_vram()
|
|
else:
|
|
clean_vram()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AILab_FinegrainBoxSegment": AILab_FinegrainBoxSegment
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AILab_FinegrainBoxSegment": "Finegrain Box Segmenter (RMBG)"
|
|
}
|