diff --git a/py/AILab_ColorToMask.py b/py/AILab_ColorToMask.py new file mode 100644 index 0000000..a094bd5 --- /dev/null +++ b/py/AILab_ColorToMask.py @@ -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) 🎭" +} diff --git a/py/AILab_Florence2.py b/py/AILab_Florence2.py new file mode 100644 index 0000000..68af3eb --- /dev/null +++ b/py/AILab_Florence2.py @@ -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": "", + "mode": "polygon", + "allows_text": True, + }, + "Phrase Grounding (text boxes)": { + "token": "", + "mode": "bbox", + "allows_text": True, + }, + "Region Proposals (boxes only)": { + "token": "", + "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)", +} diff --git a/py/AILab_Florence2Tools.py b/py/AILab_Florence2Tools.py new file mode 100644 index 0000000..27812af --- /dev/null +++ b/py/AILab_Florence2Tools.py @@ -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)", +} diff --git a/py/AILab_YoloV8.py b/py/AILab_YoloV8.py new file mode 100644 index 0000000..fe96094 --- /dev/null +++ b/py/AILab_YoloV8.py @@ -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)", +}