ComfyUI/custom_nodes/controlaltai-nodes/region_mask_generator_node.py

import torch
import numpy as np
from typing import Dict, List, Optional, Tuple
from PIL import Image, ImageDraw

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

class RegionMaskGenerator:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "width": ("INT", {"default": 1024}),
                "height": ("INT", {"default": 1024}),
                "number_of_regions": ("INT", {"default": 1, "min": 1, "max": 3}),
                # Region 1
                "region1_x1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region1_y1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region1_x2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region1_y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                # Region 2
                "region2_x1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region2_y1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region2_x2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region2_y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                # Region 3
                "region3_x1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region3_y1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region3_x2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "region3_y2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "MASK", "MASK", "INT", "BBOX", "BBOX", "BBOX")
    RETURN_NAMES = ("colored_regions_image", "bbox_preview", 
                   "mask1", "mask2", "mask3",
                   "number_of_regions",
                   "bbox1", "bbox2", "bbox3")
    FUNCTION = "generate_regions"
    CATEGORY = "ControlAltAI Nodes/Flux Region"

    def create_bbox(self, x1: float, y1: float, x2: float, y2: float) -> Dict:
        """Create bbox with debug output"""
        print(f"Creating BBOX: x1={x1:.3f}, y1={y1:.3f}, x2={x2:.3f}, y2={y2:.3f}")
        return {
            "x1": x1,
            "y1": y1,
            "x2": x2,
            "y2": y2,
            "active": True
        }

    def create_mask_from_bbox(self, bbox: Dict, width: int, height: int) -> torch.Tensor:
        """Create mask from bbox with debug output"""
        mask = torch.zeros((height, width), dtype=torch.float32)
        if bbox["active"]:
            x1 = int(bbox["x1"] * width)
            y1 = int(bbox["y1"] * height)
            x2 = int(bbox["x2"] * width)
            y2 = int(bbox["y2"] * height)
            print(f"Creating mask at pixels: x1={x1}, y1={y1}, x2={x2}, y2={y2}")
            mask[y1:y2, x1:x2] = 1.0
        return mask

    def create_preview(self, masks: List[torch.Tensor], bboxes: List[Dict], 
                      number_of_regions: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """Create preview images with debug info"""
        if not masks:
            return torch.zeros((3, 64, 64), dtype=torch.float32), torch.zeros((3, 64, 64), dtype=torch.float32)

        height, width = masks[0].shape
        
        # Create both preview images
        region_preview = Image.new("RGB", (width, height), (0, 0, 0))
        bbox_preview = Image.new("RGB", (width, height), (0, 0, 0))
        region_draw = ImageDraw.Draw(region_preview)
        bbox_draw = ImageDraw.Draw(bbox_preview)

        colors = [
            (255, 0, 0),    # Red - Region 1
            (0, 255, 0),    # Green - Region 2
            (255, 255, 0),  # Yellow - Region 3
        ]

        # Store regions for ordered preview
        preview_regions = []
        for i in range(number_of_regions):
            if bboxes[i]["active"]:
                mask_np = masks[i].cpu().numpy() > 0.5
                if mask_np.any():
                    preview_regions.append((i, mask_np, bboxes[i]))

        # Draw regions in reverse order (Region 3 first, Region 1 last)
        for i, mask_np, bbox in sorted(preview_regions, reverse=True):
            # Draw on region preview
            color_array = np.zeros((height, width, 3), dtype=np.uint8)
            for c in range(3):
                color_array[mask_np, c] = colors[i][c]
            preview_np = np.array(region_preview)
            preview_np[mask_np] = color_array[mask_np]
            region_preview = Image.fromarray(preview_np)

            # Draw on bbox preview - maintaining original bbox drawing order
            x1 = int(bbox["x1"] * width)
            y1 = int(bbox["y1"] * height)
            x2 = int(bbox["x2"] * width)
            y2 = int(bbox["y2"] * height)
            print(f"Drawing preview for region {i}: x1={x1}, y1={y1}, x2={x2}, y2={y2}")
            bbox_draw.rectangle([x1, y1, x2, y2], outline=colors[i], width=2)

        return pil2tensor(region_preview), pil2tensor(bbox_preview)

    def generate_regions(self,
                        width: int,
                        height: int,
                        number_of_regions: int,
                        **kwargs) -> Tuple:
        try:
            print(f"\nGenerating {number_of_regions} regions for {width}x{height} image")
            bboxes = []
            masks = []

            # Create regions
            for i in range(3):
                if i < number_of_regions:
                    print(f"\nProcessing region {i+1}:")
                    bbox = self.create_bbox(
                        kwargs[f"region{i+1}_x1"],
                        kwargs[f"region{i+1}_y1"],
                        kwargs[f"region{i+1}_x2"],
                        kwargs[f"region{i+1}_y2"]
                    )
                    mask = self.create_mask_from_bbox(bbox, width, height)
                    bboxes.append(bbox)
                    masks.append(mask)
                else:
                    print(f"Creating empty region {i+1}")
                    empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False}
                    bboxes.append(empty_bbox)
                    masks.append(torch.zeros((height, width), dtype=torch.float32))

            # Create previews
            region_preview, bbox_preview = self.create_preview(masks, bboxes, number_of_regions)

            return (region_preview, bbox_preview, *masks, number_of_regions, *bboxes)

        except Exception as e:
            print(f"Error in generate_regions: {str(e)}")
            empty_mask = torch.zeros((height, width), dtype=torch.float32)
            empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False}
            empty_preview = torch.zeros((3, height, width), dtype=torch.float32)
            return (empty_preview, empty_preview, 
                   empty_mask, empty_mask, empty_mask,
                   0, empty_bbox, empty_bbox, empty_bbox)

# Node class mappings
NODE_CLASS_MAPPINGS = {
    "RegionMaskGenerator": RegionMaskGenerator
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "RegionMaskGenerator": "Region Mask Generator"
}