Add files via upload
This commit is contained in:
AI Lab
@@ -0,0 +1,100 @@
|
||||
import torch
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
def color_format(color: str) -> str:
|
||||
if not color:
|
||||
return ""
|
||||
color = color.strip().upper()
|
||||
if not color.startswith('#'):
|
||||
color = f"#{color}"
|
||||
color = color[1:]
|
||||
if len(color) == 3:
|
||||
r, g, b = color[0], color[1], color[2]
|
||||
return f"#{r}{r}{g}{g}{b}{b}"
|
||||
elif len(color) < 6:
|
||||
raise ValueError(f"Invalid color format: {color}")
|
||||
elif len(color) > 6:
|
||||
color = color[:6]
|
||||
return f"#{color}"
|
||||
|
||||
def hex_to_rgb(hex_color: str):
|
||||
hex_color = hex_color.lstrip('#')
|
||||
r = int(hex_color[0:2], 16) / 255.0
|
||||
g = int(hex_color[2:4], 16) / 255.0
|
||||
b = int(hex_color[4:6], 16) / 255.0
|
||||
return r, g, b
|
||||
|
||||
def validate_image_data(images):
|
||||
if not isinstance(images, torch.Tensor):
|
||||
raise TypeError("Images must be a torch tensor.")
|
||||
if images.min() < 0 or images.max() > 255:
|
||||
raise ValueError("Images must have pixel values in the range [0, 255].")
|
||||
|
||||
def compute_color_distances(images, color):
|
||||
return torch.norm(images - color, dim=-1)
|
||||
|
||||
class AILab_ColorToMask:
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "color_to_mask"
|
||||
CATEGORY = "🧪AILab/🧽RMBG"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"invert": ("BOOLEAN", {"default": False}),
|
||||
"threshold": ("INT", {"default": 10, "min": 0, "max": 255, "step": 1}),
|
||||
"mask_color": ("COLORCODE", {"default": "#FFFFFF", "tooltip": "Mask color (hex)"}),
|
||||
},
|
||||
}
|
||||
|
||||
def color_to_mask(self, images, mask_color, threshold, invert):
|
||||
validate_image_data(images)
|
||||
|
||||
mask_color = color_format(mask_color)
|
||||
r, g, b = hex_to_rgb(mask_color)
|
||||
color = torch.tensor([r, g, b], dtype=torch.float32)
|
||||
|
||||
black = torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32)
|
||||
white = torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32)
|
||||
|
||||
if invert:
|
||||
black, white = white, black
|
||||
|
||||
per_batch = 16
|
||||
device = images.device
|
||||
color = color.to(device)
|
||||
black = black.to(device)
|
||||
white = white.to(device)
|
||||
|
||||
if images.max() > 1.0:
|
||||
images = images / 255.0
|
||||
|
||||
steps = images.shape[0]
|
||||
pbar = ProgressBar(steps)
|
||||
tensors_out = []
|
||||
|
||||
for start_idx in range(0, steps, per_batch):
|
||||
end_idx = min(start_idx + per_batch, steps)
|
||||
batch = images[start_idx:end_idx]
|
||||
|
||||
color_distances = compute_color_distances(batch, color)
|
||||
mask = color_distances <= threshold / 255.0
|
||||
mask_out = torch.where(mask.unsqueeze(-1), white, black).float()
|
||||
mask_out = mask_out.mean(dim=-1)
|
||||
|
||||
tensors_out.append(mask_out.cpu())
|
||||
pbar.update(end_idx - start_idx)
|
||||
|
||||
tensors_out = torch.cat(tensors_out, dim=0)
|
||||
tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0)
|
||||
return tensors_out,
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AILab_ColorToMask": AILab_ColorToMask
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AILab_ColorToMask": "Color to Mask (RMBG) 🎭"
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
import io
|
||||
import os
|
||||
import random
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from unittest.mock import patch
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.patches as patches
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
from PIL import Image, ImageColor, ImageDraw
|
||||
|
||||
import comfy.model_management as mm
|
||||
from comfy.utils import ProgressBar
|
||||
import folder_paths
|
||||
import transformers
|
||||
from transformers import AutoModelForCausalLM, AutoProcessor
|
||||
from transformers.dynamic_module_utils import get_imports
|
||||
|
||||
MODEL_DIR = os.path.join(folder_paths.models_dir, "LLM")
|
||||
os.makedirs(MODEL_DIR, exist_ok=True)
|
||||
folder_paths.add_model_folder_path("LLM", MODEL_DIR)
|
||||
|
||||
TASK_CONFIGS = {
|
||||
"Polygon Mask (text prompt)": {
|
||||
"token": "<REFERRING_EXPRESSION_SEGMENTATION>",
|
||||
"mode": "polygon",
|
||||
"allows_text": True,
|
||||
},
|
||||
"Phrase Grounding (text boxes)": {
|
||||
"token": "<CAPTION_TO_PHRASE_GROUNDING>",
|
||||
"mode": "bbox",
|
||||
"allows_text": True,
|
||||
},
|
||||
"Region Proposals (boxes only)": {
|
||||
"token": "<REGION_PROPOSAL>",
|
||||
"mode": "bbox",
|
||||
"allows_text": False,
|
||||
},
|
||||
}
|
||||
|
||||
TASK_CHOICES = tuple(TASK_CONFIGS.keys())
|
||||
|
||||
GENERATION_CONFIG = {
|
||||
"max_new_tokens": 512,
|
||||
"num_beams": 3,
|
||||
"do_sample": True,
|
||||
}
|
||||
|
||||
MODEL_CHOICES = (
|
||||
"microsoft/Florence-2-base",
|
||||
"microsoft/Florence-2-base-ft",
|
||||
"microsoft/Florence-2-large",
|
||||
"microsoft/Florence-2-large-ft",
|
||||
"thwri/CogFlorence-2.1-Large",
|
||||
"thwri/CogFlorence-2.2-Large",
|
||||
)
|
||||
|
||||
COLOR_BANK = ["blue", "orange", "green", "purple", "pink", "cyan"]
|
||||
|
||||
|
||||
def _fixed_get_imports(filename):
|
||||
try:
|
||||
if not str(filename).endswith("modeling_florence2.py"):
|
||||
return get_imports(filename)
|
||||
imports = get_imports(filename)
|
||||
if "flash_attn" in imports:
|
||||
imports.remove("flash_attn")
|
||||
return imports
|
||||
except Exception:
|
||||
return get_imports(filename)
|
||||
|
||||
|
||||
class AILab_Florence2:
|
||||
CATEGORY = "🧪AILab/🧽RMBG"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "JSON")
|
||||
RETURN_NAMES = ("IMAGE", "MASK", "DATA")
|
||||
FUNCTION = "analyze"
|
||||
MODEL_CACHE: Dict[Tuple[str, str, str], Dict[str, Any]] = {}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"model_name": (
|
||||
MODEL_CHOICES,
|
||||
{
|
||||
"default": MODEL_CHOICES[0],
|
||||
"tooltip": "Base = stable, +ft = fine-tuned captions, CogFlorence = sharper phrase alignment.",
|
||||
},
|
||||
),
|
||||
"task": (
|
||||
TASK_CHOICES,
|
||||
{
|
||||
"default": TASK_CHOICES[0],
|
||||
"tooltip": "Polygon masks use prompts; phrase grounding/region proposals return boxes.",
|
||||
},
|
||||
),
|
||||
"precision": (
|
||||
("fp16", "bf16", "fp32"),
|
||||
{"default": "fp16", "tooltip": "Lower precision saves VRAM; fp32 is safest if you hit NaNs."},
|
||||
),
|
||||
"attention": (
|
||||
("flash_attention_2", "sdpa", "eager"),
|
||||
{"default": "sdpa", "tooltip": "flash_attn2 needs PyTorch 2.1+; use eager if kernels fail."},
|
||||
),
|
||||
"fill_mask": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "When true, bbox tasks also output filled mask tensors."},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"output_mask_select": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Comma-separated indices or labels (e.g. 0,2,person) to limit masks.",
|
||||
},
|
||||
),
|
||||
"keep_model_loaded": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Keep weights on the current device after execution."},
|
||||
),
|
||||
"text_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"placeholder": "Prompt: e.g. a person wearing red coat",
|
||||
"tooltip": "Used for polygon masks or phrase grounding; ignored for region proposals.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def _ensure_weights(self, model_name: str) -> str:
|
||||
target = os.path.join(MODEL_DIR, model_name.split("/")[-1])
|
||||
if os.path.exists(target):
|
||||
return target
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
snapshot_download(model_name, local_dir=target, local_dir_use_symlinks=False)
|
||||
return target
|
||||
|
||||
def _get_model(self, model_name: str, precision: str, attention: str) -> Dict[str, Any]:
|
||||
key = (model_name, precision, attention)
|
||||
if key in self.MODEL_CACHE:
|
||||
return self.MODEL_CACHE[key]
|
||||
|
||||
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[precision]
|
||||
weights = self._ensure_weights(model_name)
|
||||
offload = mm.unet_offload_device()
|
||||
|
||||
if transformers.__version__ < "4.51.0":
|
||||
with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports):
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
weights,
|
||||
attn_implementation=attention,
|
||||
torch_dtype=dtype,
|
||||
trust_remote_code=True,
|
||||
).to(offload)
|
||||
else:
|
||||
from models.modeling_florence2 import Florence2ForConditionalGeneration
|
||||
|
||||
model = Florence2ForConditionalGeneration.from_pretrained(
|
||||
weights,
|
||||
attn_implementation=attention,
|
||||
torch_dtype=dtype,
|
||||
).to(offload)
|
||||
|
||||
processor = AutoProcessor.from_pretrained(weights, trust_remote_code=True)
|
||||
bundle = {"model": model, "processor": processor, "dtype": dtype}
|
||||
self.MODEL_CACHE[key] = bundle
|
||||
return bundle
|
||||
|
||||
@staticmethod
|
||||
def _prepare_prompt(task: str, text_prompt: str) -> str:
|
||||
config = TASK_CONFIGS[task]
|
||||
text = text_prompt.strip()
|
||||
if text and not config["allows_text"]:
|
||||
text = ""
|
||||
base = config["token"]
|
||||
return f"{base} {text}" if text else base
|
||||
|
||||
def _draw_regions(
|
||||
self,
|
||||
image_pil: Image.Image,
|
||||
predictions: Dict[str, Any],
|
||||
fill_mask: bool,
|
||||
select_filter: List[str],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
width, height = image_pil.size
|
||||
fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100)
|
||||
fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
|
||||
ax.imshow(image_pil)
|
||||
|
||||
mask_layer = Image.new("RGB", image_pil.size, (0, 0, 0)) if fill_mask else None
|
||||
mask_draw = ImageDraw.Draw(mask_layer) if mask_layer else None
|
||||
|
||||
for index, (bbox, label) in enumerate(zip(predictions["bboxes"], predictions["labels"])):
|
||||
x0, y0, x1, y1 = bbox
|
||||
if y1 < y0:
|
||||
y0, y1 = y1, y0
|
||||
if x1 < x0:
|
||||
x0, x1 = x1, x0
|
||||
filter_hit = not select_filter or str(index) in select_filter or label in select_filter
|
||||
if fill_mask and filter_hit:
|
||||
mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255))
|
||||
rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=1, edgecolor="red", facecolor="none")
|
||||
ax.add_patch(rect)
|
||||
ax.text(
|
||||
x0,
|
||||
max(0, y0 - 12),
|
||||
f"{index}.{label}",
|
||||
color="white",
|
||||
fontsize=12,
|
||||
bbox=dict(facecolor=random.choice(COLOR_BANK), alpha=0.5),
|
||||
)
|
||||
|
||||
ax.axis("off")
|
||||
buf = io.BytesIO()
|
||||
plt.savefig(buf, format="png", pad_inches=0)
|
||||
buf.seek(0)
|
||||
annotated = Image.open(buf)
|
||||
plt.close(fig)
|
||||
|
||||
annotated_tensor = TF.to_tensor(annotated)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
||||
if mask_layer is not None:
|
||||
mask_tensor = TF.to_tensor(mask_layer).unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
||||
mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0]
|
||||
else:
|
||||
mask_tensor = torch.zeros((1, annotated_tensor.shape[1], annotated_tensor.shape[2]), dtype=torch.float32)
|
||||
|
||||
return annotated_tensor, mask_tensor
|
||||
|
||||
def _segment_polygons(
|
||||
self,
|
||||
image_pil: Image.Image,
|
||||
predictions: Dict[str, Any],
|
||||
fill_mask: bool,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
width, height = image_pil.size
|
||||
mask_image = Image.new("RGB", (width, height), "black")
|
||||
mask_draw = ImageDraw.Draw(mask_image)
|
||||
|
||||
for polygons, label in zip(predictions["polygons"], predictions["labels"]):
|
||||
color = random.choice(COLOR_BANK)
|
||||
for polygon in polygons:
|
||||
polygon = np.array(polygon).reshape(-1, 2)
|
||||
polygon = np.clip(polygon, [0, 0], [width - 1, height - 1])
|
||||
if len(polygon) < 3:
|
||||
continue
|
||||
pts = polygon.reshape(-1).tolist()
|
||||
if fill_mask:
|
||||
overlay = Image.new("RGBA", image_pil.size, (255, 255, 255, 0))
|
||||
draw = ImageDraw.Draw(overlay)
|
||||
rgba = ImageColor.getrgb(color) + (180,)
|
||||
draw.polygon(pts, outline=color, fill=rgba, width=3)
|
||||
image_pil = Image.alpha_composite(image_pil.convert("RGBA"), overlay).convert("RGB")
|
||||
else:
|
||||
draw = ImageDraw.Draw(image_pil)
|
||||
draw.polygon(pts, outline=color, width=3)
|
||||
mask_draw.polygon(pts, outline="white", fill="white")
|
||||
|
||||
image_tensor = TF.to_tensor(image_pil)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
||||
mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
||||
mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0]
|
||||
return image_tensor, mask_tensor
|
||||
|
||||
def analyze(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
model_name: str,
|
||||
task: str,
|
||||
precision: str,
|
||||
attention: str,
|
||||
fill_mask: bool,
|
||||
output_mask_select: str = "",
|
||||
keep_model_loaded: bool = False,
|
||||
text_prompt: str = "",
|
||||
):
|
||||
bundle = self._get_model(model_name, precision, attention)
|
||||
model = bundle["model"]
|
||||
processor = bundle["processor"]
|
||||
dtype = bundle["dtype"]
|
||||
|
||||
device = mm.get_torch_device()
|
||||
offload = mm.unet_offload_device()
|
||||
model.to(device)
|
||||
images = image.permute(0, 3, 1, 2)
|
||||
|
||||
task_config = TASK_CONFIGS[task]
|
||||
prompt = self._prepare_prompt(task, text_prompt)
|
||||
out_images, out_masks, out_data = [], [], []
|
||||
pbar = ProgressBar(len(images))
|
||||
select_filter = [s.strip() for s in output_mask_select.split(",") if s.strip()]
|
||||
|
||||
for tensor in images:
|
||||
image_pil = TF.to_pil_image(tensor)
|
||||
inputs = processor(
|
||||
text=prompt,
|
||||
images=image_pil,
|
||||
return_tensors="pt",
|
||||
do_rescale=False,
|
||||
).to(dtype).to(device)
|
||||
|
||||
generated_ids = model.generate(
|
||||
input_ids=inputs["input_ids"],
|
||||
pixel_values=inputs["pixel_values"],
|
||||
max_new_tokens=GENERATION_CONFIG["max_new_tokens"],
|
||||
do_sample=GENERATION_CONFIG["do_sample"],
|
||||
num_beams=GENERATION_CONFIG["num_beams"],
|
||||
use_cache=False,
|
||||
)
|
||||
raw = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
||||
parsed = processor.post_process_generation(raw, task=task_config["token"], image_size=image_pil.size)
|
||||
predictions = parsed[task_config["token"]]
|
||||
|
||||
if task_config["mode"] == "bbox":
|
||||
annotated, mask = self._draw_regions(image_pil, predictions, fill_mask, select_filter)
|
||||
out_images.append(annotated)
|
||||
out_masks.append(mask)
|
||||
else:
|
||||
image_tensor, mask_tensor = self._segment_polygons(image_pil, predictions, fill_mask)
|
||||
out_images.append(image_tensor)
|
||||
out_masks.append(mask_tensor)
|
||||
|
||||
out_data.append(predictions)
|
||||
pbar.update(1)
|
||||
|
||||
image_out = torch.cat(out_images, dim=0)
|
||||
mask_out = torch.cat(out_masks, dim=0)
|
||||
|
||||
if not keep_model_loaded:
|
||||
model.to(offload)
|
||||
mm.soft_empty_cache()
|
||||
|
||||
return image_out, mask_out, out_data
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AILab_Florence2": AILab_Florence2,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AILab_Florence2": "Florence2 (RMBG)",
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
import json
|
||||
from typing import Any, List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class AILab_Florence2ToCoordinates:
|
||||
CATEGORY = "🧪AILab/🧽RMBG"
|
||||
RETURN_TYPES = ("STRING", "BBOX", "MASK")
|
||||
RETURN_NAMES = ("CENTER_COORDINATES", "BBOXES", "MASK")
|
||||
FUNCTION = "convert"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"data": ("JSON", {"tooltip": "Florence2 JSON output (list per image)."}),
|
||||
"index": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Comma-separated indexes; blank = use all boxes from first item.",
|
||||
},
|
||||
),
|
||||
"batch": ("BOOLEAN", {"default": False, "tooltip": "If true, gather boxes across the batch."}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",)
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _parse_payload(payload: Any) -> List[Any]:
|
||||
if payload is None:
|
||||
return []
|
||||
if isinstance(payload, str):
|
||||
cleaned = payload.strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
try:
|
||||
return json.loads(cleaned)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback: some upstream nodes stringify Python lists with single quotes
|
||||
try:
|
||||
normalized = cleaned.replace("'", '"')
|
||||
return json.loads(normalized)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError("Invalid JSON payload for Florence2 data") from exc
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _get_bboxes(entry: Any) -> List[List[float]]:
|
||||
if isinstance(entry, dict):
|
||||
if "bboxes" in entry:
|
||||
return entry["bboxes"]
|
||||
raise ValueError("Entry does not contain 'bboxes'.")
|
||||
if isinstance(entry, list):
|
||||
return entry
|
||||
raise ValueError("Unsupported entry type; expected dict with 'bboxes' or list of boxes.")
|
||||
|
||||
@staticmethod
|
||||
def _parse_indexes(index_str: str, default_count: int) -> List[int]:
|
||||
text = index_str.strip()
|
||||
if not text:
|
||||
return list(range(default_count))
|
||||
try:
|
||||
return [int(part.strip()) for part in text.split(",") if part.strip()]
|
||||
except ValueError as exc:
|
||||
raise ValueError("Index must be comma-separated integers.") from exc
|
||||
|
||||
def convert(self, data, index: str, batch: bool = False, image=None):
|
||||
records = self._parse_payload(data)
|
||||
if not records:
|
||||
empty = json.dumps([{"x": 0, "y": 0}])
|
||||
mask = self._build_empty_mask(image)
|
||||
return (empty, [], mask)
|
||||
|
||||
first_bboxes = self._get_bboxes(records[0])
|
||||
if not first_bboxes:
|
||||
empty = json.dumps([{"x": 0, "y": 0}])
|
||||
mask = self._build_empty_mask(image)
|
||||
return (empty, [], mask)
|
||||
|
||||
indexes = self._parse_indexes(index, len(first_bboxes))
|
||||
centers = []
|
||||
selected_boxes = []
|
||||
selections: List[Tuple[int, List[float]]] = []
|
||||
max_dims = {}
|
||||
|
||||
def append_box(batch_idx: int, box: List[float]):
|
||||
min_x, min_y, max_x, max_y = box
|
||||
center_x = int((min_x + max_x) / 2)
|
||||
center_y = int((min_y + max_y) / 2)
|
||||
centers.append({"x": center_x, "y": center_y})
|
||||
selected_boxes.append(box)
|
||||
selections.append((batch_idx, box))
|
||||
dims = max_dims.setdefault(batch_idx, [1, 1])
|
||||
dims[0] = max(dims[0], int(max_x) + 1)
|
||||
dims[1] = max(dims[1], int(max_y) + 1)
|
||||
|
||||
if batch:
|
||||
for batch_idx, record in enumerate(records):
|
||||
boxes = self._get_bboxes(record)
|
||||
for idx in indexes:
|
||||
if 0 <= idx < len(boxes):
|
||||
append_box(batch_idx, boxes[idx])
|
||||
else:
|
||||
boxes = first_bboxes
|
||||
for idx in indexes:
|
||||
if not 0 <= idx < len(boxes):
|
||||
raise ValueError(f"Index {idx} is out of range for available boxes")
|
||||
append_box(0, boxes[idx])
|
||||
|
||||
mask_tensor = self._build_mask_tensor(image, selections, max_dims, batch, len(records))
|
||||
|
||||
return (json.dumps(centers), selected_boxes, mask_tensor)
|
||||
|
||||
@staticmethod
|
||||
def _build_empty_mask(image):
|
||||
if image is not None:
|
||||
tensor = image
|
||||
if tensor.dim() == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
return torch.zeros((tensor.shape[0], tensor.shape[1], tensor.shape[2]), dtype=torch.float32, device=tensor.device)
|
||||
return torch.zeros((1, 1, 1), dtype=torch.float32)
|
||||
|
||||
def _build_mask_tensor(self, image, selections, max_dims, batch_mode, record_count):
|
||||
if not selections:
|
||||
return self._build_empty_mask(image)
|
||||
|
||||
if image is not None:
|
||||
tensor = image
|
||||
if tensor.dim() == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
base = torch.zeros((tensor.shape[0], tensor.shape[1], tensor.shape[2]), dtype=torch.float32, device=tensor.device)
|
||||
else:
|
||||
batch_size = record_count if batch_mode else 1
|
||||
max_width = max((dims[0] for dims in max_dims.values()), default=1)
|
||||
max_height = max((dims[1] for dims in max_dims.values()), default=1)
|
||||
base = torch.zeros((batch_size, max_height, max_width), dtype=torch.float32)
|
||||
|
||||
for batch_idx, box in selections:
|
||||
min_x, min_y, max_x, max_y = box
|
||||
width = base.shape[2]
|
||||
height = base.shape[1]
|
||||
x0 = max(0, min(int(min_x), width - 1))
|
||||
y0 = max(0, min(int(min_y), height - 1))
|
||||
x1 = max(x0 + 1, min(int(max_x), width))
|
||||
y1 = max(y0 + 1, min(int(max_y), height))
|
||||
base[batch_idx, y0:y1, x0:x1] = 1.0
|
||||
|
||||
return base
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AILab_Florence2ToCoordinates": AILab_Florence2ToCoordinates,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AILab_Florence2ToCoordinates": "Florence2 Box Coordinates (RMBG)",
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
import os
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
import folder_paths
|
||||
from AILab_ImageMaskTools import pil2tensor, tensor2pil
|
||||
|
||||
ULTRALYTICS_DIR = os.path.join(folder_paths.models_dir, "ultralytics")
|
||||
YOLO_LEGACY_DIR = os.path.join(folder_paths.models_dir, "yolo")
|
||||
os.makedirs(ULTRALYTICS_DIR, exist_ok=True)
|
||||
os.makedirs(YOLO_LEGACY_DIR, exist_ok=True)
|
||||
folder_paths.add_model_folder_path("ultralytics", ULTRALYTICS_DIR, is_default=True)
|
||||
folder_paths.add_model_folder_path("ultralytics", YOLO_LEGACY_DIR)
|
||||
|
||||
DEVICE_CHOICES = ("auto", "cuda", "cpu", "mps")
|
||||
MASK_COUNT_CHOICES = ("all", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10")
|
||||
MASK_INDEX_CHOICES = ("none", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10")
|
||||
|
||||
|
||||
class AILab_YoloV8Adv:
|
||||
CATEGORY = "🧪AILab/🧽RMBG"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "MASK")
|
||||
RETURN_NAMES = ("ANNOTATED_IMAGE", "MASK", "MASK_LIST")
|
||||
FUNCTION = "yolo_detect"
|
||||
|
||||
_MODEL_CACHE: Dict[str, "YOLO"] = {}
|
||||
|
||||
@classmethod
|
||||
def _list_models(cls) -> List[str]:
|
||||
files = folder_paths.get_filename_list("ultralytics")
|
||||
return sorted(f for f in files if f.lower().endswith(".pt"))
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
models = cls._list_models()
|
||||
if not models:
|
||||
models = [f"Put .pt models into {ULTRALYTICS_DIR}"]
|
||||
default_model = models[0]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR} (subfolders allowed)."}),
|
||||
"mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}),
|
||||
},
|
||||
"optional": {
|
||||
"select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}),
|
||||
"conf": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Confidence threshold forwarded to Ultralytics."}),
|
||||
"iou": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "IOU used during NMS."}),
|
||||
"classes": ("STRING", {"default": "", "placeholder": "e.g. 0,2,5-7", "tooltip": "Comma list or ranges of class IDs; empty keeps every class."}),
|
||||
"device": (DEVICE_CHOICES, {"default": "auto", "tooltip": "Force a device or auto-detect CUDA → MPS → CPU."}),
|
||||
"max_det": ("INT", {"default": 300, "min": 1, "max": 1000, "step": 1, "tooltip": "Maximum detections per image."}),
|
||||
"retina_masks": ("BOOLEAN", {"default": True, "tooltip": "Use high-resolution masks (Ultralytics retina_masks flag)."}),
|
||||
"agnostic_nms": ("BOOLEAN", {"default": False, "tooltip": "Enable class-agnostic NMS."}),
|
||||
},
|
||||
}
|
||||
|
||||
def _resolve_device(self, requested: str) -> str:
|
||||
if requested != "auto":
|
||||
return requested
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
|
||||
def _parse_classes(self, value: str) -> Optional[List[int]]:
|
||||
if not value or not value.strip():
|
||||
return None
|
||||
classes: List[int] = []
|
||||
try:
|
||||
for chunk in value.split(","):
|
||||
chunk = chunk.strip()
|
||||
if not chunk:
|
||||
continue
|
||||
if "-" in chunk:
|
||||
start, end = [int(x) for x in chunk.split("-", 1)]
|
||||
if start > end:
|
||||
start, end = end, start
|
||||
classes.extend(range(start, end + 1))
|
||||
else:
|
||||
classes.append(int(chunk))
|
||||
return sorted(set(classes))
|
||||
except ValueError:
|
||||
print(f"[AILab_YoloV8] Invalid classes string: {value}. Ignoring filter.")
|
||||
return None
|
||||
|
||||
def _resolve_model_path(self, name: str) -> str:
|
||||
return folder_paths.get_full_path_or_raise("ultralytics", name)
|
||||
|
||||
def _get_model(self, model_path: str):
|
||||
model = self._MODEL_CACHE.get(model_path)
|
||||
if model is None:
|
||||
from ultralytics import YOLO
|
||||
model = YOLO(model_path)
|
||||
self._MODEL_CACHE[model_path] = model
|
||||
return model
|
||||
|
||||
def _result_to_tensor(self, result) -> torch.Tensor:
|
||||
plotted = result.plot()
|
||||
rgb = plotted[..., ::-1]
|
||||
return pil2tensor(Image.fromarray(rgb))
|
||||
|
||||
def _mask_from_tensor(self, mask_tensor: torch.Tensor, size: Image.Image.size):
|
||||
mask_np = mask_tensor.detach().cpu().numpy()
|
||||
mask_img = Image.fromarray((mask_np * 255).astype(np.uint8))
|
||||
if mask_img.size != size:
|
||||
mask_img = mask_img.resize(size, Image.Resampling.NEAREST)
|
||||
return torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0)
|
||||
|
||||
def _collect_masks(self, result, size) -> List[torch.Tensor]:
|
||||
width, height = size
|
||||
masks: List[torch.Tensor] = []
|
||||
|
||||
if getattr(result, "masks", None) is not None and result.masks.data is not None:
|
||||
for mask_tensor in result.masks.data:
|
||||
masks.append(self._mask_from_tensor(mask_tensor, size))
|
||||
|
||||
elif getattr(result, "boxes", None) is not None and len(result.boxes.xyxy) > 0:
|
||||
for box in result.boxes:
|
||||
x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]
|
||||
mask_img = Image.new("L", size, 0)
|
||||
draw = ImageDraw.Draw(mask_img)
|
||||
draw.rectangle([x1, y1, x2, y2], fill=255)
|
||||
masks.append(torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0))
|
||||
|
||||
if not masks:
|
||||
masks.append(torch.zeros((height, width), dtype=torch.float32))
|
||||
|
||||
return masks
|
||||
|
||||
def _merge_masks(self, masks: List[torch.Tensor]) -> torch.Tensor:
|
||||
if not masks:
|
||||
raise ValueError("Cannot merge empty mask list.")
|
||||
merged = torch.zeros_like(masks[0])
|
||||
|
||||
for mask in masks:
|
||||
merged = torch.maximum(merged, mask)
|
||||
|
||||
return merged
|
||||
|
||||
def yolo_detect(
|
||||
self,
|
||||
images,
|
||||
yolo_model,
|
||||
mask_count="all",
|
||||
conf=0.25,
|
||||
iou=0.45,
|
||||
classes="",
|
||||
device="auto",
|
||||
max_det=300,
|
||||
retina_masks=True,
|
||||
agnostic_nms=False,
|
||||
select_mask_index: str = "none",
|
||||
):
|
||||
model_path = self._resolve_model_path(yolo_model)
|
||||
model = self._get_model(model_path)
|
||||
device_target = self._resolve_device(device)
|
||||
class_filter = self._parse_classes(classes)
|
||||
|
||||
merged_masks: List[torch.Tensor] = []
|
||||
annotated_images: List[torch.Tensor] = []
|
||||
mask_list: List[torch.Tensor] = []
|
||||
|
||||
count_limit = 0 if mask_count == "all" else max(0, int(mask_count))
|
||||
chosen_index: Optional[int] = None
|
||||
if select_mask_index != "none":
|
||||
chosen_index = int(select_mask_index) - 1
|
||||
|
||||
for idx in range(images.shape[0]):
|
||||
image_pil = tensor2pil(images[idx])
|
||||
|
||||
results = model(
|
||||
image_pil,
|
||||
conf=conf,
|
||||
iou=iou,
|
||||
classes=class_filter,
|
||||
device=device_target,
|
||||
max_det=max_det,
|
||||
retina_masks=retina_masks,
|
||||
agnostic_nms=agnostic_nms,
|
||||
)
|
||||
|
||||
if not results:
|
||||
continue
|
||||
|
||||
result = results[0]
|
||||
annotated_images.append(self._result_to_tensor(result))
|
||||
|
||||
frame_masks = self._collect_masks(result, image_pil.size)
|
||||
|
||||
selected_masks: List[torch.Tensor]
|
||||
if chosen_index is None:
|
||||
if count_limit <= 0 or count_limit >= len(frame_masks):
|
||||
selected_masks = frame_masks
|
||||
else:
|
||||
selected_masks = frame_masks[:count_limit]
|
||||
else:
|
||||
if chosen_index >= len(frame_masks):
|
||||
selected_masks = []
|
||||
else:
|
||||
span = count_limit if count_limit > 0 else 1
|
||||
selected_masks = frame_masks[chosen_index : chosen_index + span]
|
||||
|
||||
if selected_masks:
|
||||
merged_masks.append(self._merge_masks(selected_masks))
|
||||
mask_list.extend(selected_masks)
|
||||
else:
|
||||
fallback = torch.zeros_like(frame_masks[0])
|
||||
merged_masks.append(fallback)
|
||||
mask_list.append(fallback)
|
||||
|
||||
if not merged_masks:
|
||||
width, height = tensor2pil(images[0]).size
|
||||
merged_masks = [torch.zeros((height, width), dtype=torch.float32)]
|
||||
|
||||
if not mask_list:
|
||||
width, height = merged_masks[0].shape[1], merged_masks[0].shape[0]
|
||||
mask_list = [torch.zeros((height, width), dtype=torch.float32)]
|
||||
|
||||
if not annotated_images:
|
||||
annotated_images = [images]
|
||||
|
||||
merged_tensor = torch.stack(merged_masks, dim=0)
|
||||
annotated_tensor = torch.cat(annotated_images, dim=0)
|
||||
mask_tensor = torch.stack(mask_list, dim=0)
|
||||
|
||||
return annotated_tensor, merged_tensor, mask_tensor
|
||||
|
||||
|
||||
class AILab_YoloV8(AILab_YoloV8Adv):
|
||||
FUNCTION = "yolo_detect_simple"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
models = cls._list_models()
|
||||
if not models:
|
||||
models = [f"Put .pt models into {ULTRALYTICS_DIR}"]
|
||||
default_model = models[0]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR}. Advanced controls available on YOLOv8 Adv."}),
|
||||
"mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}),
|
||||
},
|
||||
"optional": {
|
||||
"select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}),
|
||||
},
|
||||
}
|
||||
|
||||
def yolo_detect_simple(self, images, yolo_model, mask_count="all", select_mask_index="none"):
|
||||
return super().yolo_detect(
|
||||
images=images,
|
||||
yolo_model=yolo_model,
|
||||
mask_count=mask_count,
|
||||
conf=0.25,
|
||||
iou=0.45,
|
||||
classes="",
|
||||
device="auto",
|
||||
max_det=300,
|
||||
retina_masks=True,
|
||||
agnostic_nms=False,
|
||||
select_mask_index=select_mask_index,
|
||||
)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AILab_YoloV8": AILab_YoloV8,
|
||||
"AILab_YoloV8Adv": AILab_YoloV8Adv,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AILab_YoloV8": "YOLOv8 (RMBG)",
|
||||
"AILab_YoloV8Adv": "YOLOv8 Adv (RMBG)",
|
||||
}
|
||||
Reference in New Issue
Block a user