diff --git a/README.md b/README.md index 6f563f8..125aa79 100644 --- a/README.md +++ b/README.md @@ -62,17 +62,52 @@ Inputs: - `mask1`: Optional first mask group. Output order is anchored to this group when both inputs are connected. - `mask2`: Optional second mask group to match against `mask1`. +- `output_mode`: `mask` keeps the original product mask shape inside each bbox slice. `bbox` outputs full rectangular bbox slices for the most stable spatial correspondence. Default is `bbox`. - `min_area_ratio`: Removes tiny disconnected components before grouping and matching, with `0.01` step precision. Default is `0.01`; `0.0` disables filtering. Values are relative to the largest disconnected component in the same input group, so `0.02` removes components smaller than 2% of that largest component's foreground area. - `rows`: Number of equal-height bbox grid rows. Default is `4`. - `columns`: Number of equal-width bbox grid columns. Default is `1`. -- `output_mode`: `mask` keeps the original product mask shape inside each bbox slice. `bbox` outputs full rectangular bbox slices for the most stable spatial correspondence. Default is `bbox`. +- `auto_match_orientation`: Automatically swaps `rows` and `columns` for horizontal `mask1` product bboxes. Default is enabled. Outputs: - `mask1`: First sliced/grouped mask batch. - `mask2`: Second mask batch aligned to `mask1`. When only one input is connected, this is an empty placeholder batch. -The node splits disconnected painted islands internally, removes components below `min_area_ratio`, unions the remaining regions into one product mask, and slices that product bbox into a row-major grid: top to bottom, left to right. In `mask` mode, empty `mask1` cells are skipped and `mask2` keeps the same remaining cell positions. In `bbox` mode, the full `rows × columns` rectangular grid is preserved. Use `rows=4, columns=1` for vertical product labels, `rows=1, columns=4` for horizontal labels, and multiple rows and columns for two-dimensional local regions. +The node splits disconnected painted islands internally, removes components below `min_area_ratio`, unions the remaining regions into one product mask, and slices that product bbox into a row-major grid: top to bottom, left to right. When `auto_match_orientation` is enabled, `rows=4, columns=1` stays 4x1 for vertical products and becomes 1x4 for horizontal products. In `mask` mode, empty `mask1` cells are skipped and `mask2` keeps the same remaining cell positions. In `bbox` mode, the full `rows × columns` rectangular grid is preserved. Use multiple rows and columns for two-dimensional local regions. + +### RefineNode Match Product Angle + +Rotates one product image to match the main mask angle of a reference product mask. The node rotates the whole source image without scaling, rotates the source mask with the same angle, and crops the result to the rotated source mask bbox. + +Inputs: + +- `image`: Source product image to rotate. +- `source_mask`: Product mask for the source image. +- `reference_mask`: Product mask whose angle should be matched. +- `canvas_expand`: Expands the output crop from the rotated product bbox toward the full rotated canvas. `0.0` keeps the tight product bbox crop; `1.0` keeps the full rotated canvas. + +Outputs: + +- `image`: Rotated RGB image cropped to the rotated source mask bbox. +- `mask`: Rotated source mask cropped to the same bbox. +- `angle`: Applied rotation angle in degrees. + +The angle is estimated from each mask's foreground principal axis. Empty source or reference masks return the original source image, the source mask, and `angle=0`. The node does not scale, perform perspective correction, or distinguish 0 degrees from 180 degrees. + +### RefineNode Rotate Image + +Rotates an image by a manually specified angle. + +Inputs: + +- `image`: Image to rotate. +- `angle`: Rotation angle in degrees. + +Outputs: + +- `image`: Rotated RGB image. + +The image is rotated without scaling. The output canvas expands to contain the full rotated image, and newly exposed pixels are filled with black. ### RefineNode Preprocess Mask diff --git a/nodes.py b/nodes.py index 93db377..73611fa 100644 --- a/nodes.py +++ b/nodes.py @@ -215,8 +215,10 @@ def tensor_image_to_pil(images: torch.Tensor, index: int) -> Image.Image: raise ValueError(f"Expected IMAGE tensor with 3 or 4 dims, got {tuple(images.shape)}") if tensor.shape[-1] == 1: tensor = tensor.repeat(1, 1, 3) + elif tensor.shape[-1] == 4: + tensor = tensor[..., :3] if tensor.shape[-1] != 3: - raise ValueError(f"Expected IMAGE tensor channel count 3, got {tensor.shape[-1]}") + raise ValueError(f"Expected IMAGE tensor channel count 1, 3, or 4, got {tensor.shape[-1]}") arr = (tensor.clamp(0, 1).numpy() * 255.0 + 0.5).astype(np.uint8) return Image.fromarray(arr, mode="RGB") @@ -509,6 +511,7 @@ def slice_masks_by_product_bbox( min_area_ratio: float, rows: int, columns: int, + auto_match_orientation: bool, output_mode: str, ) -> tuple[list[Image.Image], list[Image.Image]]: mask1_union = filtered_union_mask(mask1_images, min_area_ratio) @@ -519,6 +522,9 @@ def slice_masks_by_product_bbox( if mask1_bbox is None: return ([Image.new("L", mask1_size, 0)], [Image.new("L", mask2_size, 0)]) + if auto_match_orientation and (mask1_bbox[2] - mask1_bbox[0]) > (mask1_bbox[3] - mask1_bbox[1]): + rows, columns = columns, rows + cells = normalized_grid_cells(rows, columns) mask1_slices = bbox_grid_sliced_masks(mask1_union, mask1_bbox, cells, output_mode) @@ -621,6 +627,89 @@ def tensor_mask_as_list_item(mask_l: Image.Image) -> torch.Tensor: return pil_to_tensor_mask(mask_l).unsqueeze(0) +def normalize_axis_angle(angle: float) -> float: + while angle < -90.0: + angle += 180.0 + while angle >= 90.0: + angle -= 180.0 + return angle + + +def mask_principal_axis_angle(mask_l: Image.Image) -> float | None: + arr = np.asarray(mask_l.convert("L"), dtype=np.uint8) > 0 + ys, xs = np.nonzero(arr) + if xs.size < 2: + return None + coords = np.stack([xs.astype(np.float64), -ys.astype(np.float64)], axis=1) + coords -= coords.mean(axis=0, keepdims=True) + cov = np.cov(coords, rowvar=False) + if not np.isfinite(cov).all(): + return None + values, vectors = np.linalg.eigh(cov) + vector = vectors[:, int(np.argmax(values))] + if not np.isfinite(vector).all() or np.linalg.norm(vector) <= 1e-8: + return None + return normalize_axis_angle(math.degrees(math.atan2(float(vector[1]), float(vector[0])))) + + +def rotated_product_crop( + image: Image.Image, + mask_l: Image.Image, + angle: float, + canvas_expand: float = 0.0, +) -> tuple[Image.Image, Image.Image]: + rotated_image = image.convert("RGB").rotate( + angle, + resample=_BICUBIC, + expand=True, + fillcolor=(0, 0, 0), + ) + rotated_mask = binary_mask(mask_l).rotate( + angle, + resample=_NEAREST, + expand=True, + fillcolor=0, + ) + rotated_mask = binary_mask(rotated_mask) + bbox = bbox_from_mask_or_none(rotated_mask) + if bbox is None: + return rotated_image, rotated_mask + expand = max(0.0, min(1.0, float(canvas_expand))) + if expand <= 0.0: + crop_box = bbox + else: + canvas_w, canvas_h = rotated_image.size + x1, y1, x2, y2 = bbox + crop_box = ( + max(0, int(math.floor(x1 * (1.0 - expand)))), + max(0, int(math.floor(y1 * (1.0 - expand)))), + min(canvas_w, int(math.ceil(x2 * (1.0 - expand) + canvas_w * expand))), + min(canvas_h, int(math.ceil(y2 * (1.0 - expand) + canvas_h * expand))), + ) + return rotated_image.crop(crop_box), rotated_mask.crop(crop_box) + + +def rotate_image_canvas(image: Image.Image, angle: float) -> Image.Image: + return image.convert("RGB").rotate( + float(angle), + resample=_BICUBIC, + expand=True, + fillcolor=(0, 0, 0), + ) + + +def stack_image_pils(images: list[Image.Image]) -> torch.Tensor: + if not images: + raise ValueError("Cannot output an empty image batch.") + size = images[0].size + if any(image.size != size for image in images): + raise ValueError( + "RefineNode Match Product Angle batch outputs have different crop sizes. " + "Split the batch or use inputs with matching rotated product bbox sizes." + ) + return torch.stack([pil_to_tensor_image(image) for image in images], dim=0) + + def upscale_image(samples: torch.Tensor, width: int, height: int, method: str, crop: str = "disabled") -> torch.Tensor: import comfy.utils @@ -932,10 +1021,11 @@ class RefineNodeSliceAndMatchMasks: def INPUT_TYPES(cls): return { "required": { + "output_mode": (["mask", "bbox"], {"default": "bbox"}), "min_area_ratio": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.01}), "rows": ("INT", {"default": 4, "min": 1, "max": 16}), "columns": ("INT", {"default": 1, "min": 1, "max": 16}), - "output_mode": (["mask", "bbox"], {"default": "bbox"}), + "auto_match_orientation": ("BOOLEAN", {"default": True}), }, "optional": { "mask1": ("MASK",), @@ -954,6 +1044,7 @@ class RefineNodeSliceAndMatchMasks: min_area_ratio: float | list[float], rows: int | list[int], columns: int | list[int], + auto_match_orientation: bool | list[bool], output_mode: str | list[str], mask1: torch.Tensor | list[torch.Tensor] | None = None, mask2: torch.Tensor | list[torch.Tensor] | None = None, @@ -961,6 +1052,7 @@ class RefineNodeSliceAndMatchMasks: min_area_ratio = clamp_unit_float(min_area_ratio, 0.01) rows = clamp_int_value(rows, 4, 1, 16) columns = clamp_int_value(columns, 1, 1, 16) + auto_match_orientation = bool(auto_match_orientation[0]) if isinstance(auto_match_orientation, list) else bool(auto_match_orientation) output_mode = normalize_choice(output_mode, "bbox", {"mask", "bbox"}) if mask1 is None and mask2 is None: raise ValueError("Connect at least one mask input.") @@ -975,11 +1067,100 @@ class RefineNodeSliceAndMatchMasks: min_area_ratio, rows, columns, + auto_match_orientation, output_mode, ) return (stack_mask_images(sliced_mask1), stack_mask_images(sliced_mask2)) +class RefineNodeMatchProductAngle: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "source_mask": ("MASK",), + "reference_mask": ("MASK",), + "canvas_expand": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK", "FLOAT") + RETURN_NAMES = ("image", "mask", "angle") + OUTPUT_IS_LIST = (False, False, True) + FUNCTION = "match_angle" + CATEGORY = "RefineNode/Transform" + + def match_angle( + self, + image: torch.Tensor, + source_mask: torch.Tensor, + reference_mask: torch.Tensor, + canvas_expand: float | int | list[float] | list[int], + ): + canvas_expand = clamp_unit_float(canvas_expand, 0.0) + batch = max(image_batch_size(image), mask_batch_size(source_mask), mask_batch_size(reference_mask)) + reference_size = mask_tensor_spatial_size(reference_mask) + output_images: list[Image.Image] = [] + output_masks: list[Image.Image] = [] + angles: list[float] = [] + + for index in range(batch): + source_image = normalize_to_srgb(tensor_image_to_pil(image, index)) + source_mask_l = tensor_mask_to_pil(source_mask, index, source_image.size) + reference_mask_l = tensor_mask_to_pil(reference_mask, index, reference_size) + + source_angle = mask_principal_axis_angle(source_mask_l) + reference_angle = mask_principal_axis_angle(reference_mask_l) + if source_angle is None or reference_angle is None: + output_images.append(source_image) + output_masks.append(binary_mask(source_mask_l)) + angles.append(0.0) + continue + + angle = normalize_axis_angle(reference_angle - source_angle) + rotated_image, rotated_mask = rotated_product_crop(source_image, source_mask_l, angle, canvas_expand) + output_images.append(rotated_image) + output_masks.append(rotated_mask) + angles.append(float(angle)) + + return (stack_image_pils(output_images), stack_mask_images(output_masks), angles) + + +class RefineNodeRotateImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "angle": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "rotate" + CATEGORY = "RefineNode/Transform" + + def rotate( + self, + image: torch.Tensor, + angle: float | int | list[float] | list[int], + ): + if isinstance(angle, list): + angle = angle[0] if angle else 0.0 + try: + angle = float(angle) + except (TypeError, ValueError): + angle = 0.0 + + output_images = [ + rotate_image_canvas(normalize_to_srgb(tensor_image_to_pil(image, index)), angle) + for index in range(image_batch_size(image)) + ] + return (stack_image_pils(output_images),) + + class RefineNodePreprocessMask: @classmethod def INPUT_TYPES(cls): @@ -1200,8 +1381,10 @@ class RefineNodeReferenceImageProcess: raise ValueError(f"Expected IMAGE tensor with 3 or 4 dims, got {tuple(image.shape)}") if sample.shape[-1] == 1: sample = sample.repeat(1, 1, 1, 3) + elif sample.shape[-1] == 4: + sample = sample[..., :3] if sample.shape[-1] != 3: - raise ValueError(f"Expected IMAGE tensor channel count 3, got {sample.shape[-1]}") + raise ValueError(f"Expected IMAGE tensor channel count 1, 3, or 4, got {sample.shape[-1]}") sample = sample.movedim(-1, 1) out = upscale_to_kontext_size(sample, width, height, resize_method, crop_mode).movedim(1, -1) processed.append(out[0]) @@ -1430,6 +1613,8 @@ def focus_crop_region( NODE_CLASS_MAPPINGS = { "RefineNodeMaskBatchProcess": RefineNodeMaskBatchProcess, "RefineNodeSliceAndMatchMasks": RefineNodeSliceAndMatchMasks, + "RefineNodeMatchProductAngle": RefineNodeMatchProductAngle, + "RefineNodeRotateImage": RefineNodeRotateImage, "RefineNodePreprocessMask": RefineNodePreprocessMask, "RefineNodeReferenceImageProcess": RefineNodeReferenceImageProcess, "RefineNodeRestoreMaskToOriginal": RefineNodeRestoreMaskToOriginal, @@ -1439,6 +1624,8 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "RefineNodeMaskBatchProcess": "RefineNode Mask Batch Process", "RefineNodeSliceAndMatchMasks": "RefineNode Slice And Match Masks", + "RefineNodeMatchProductAngle": "RefineNode Match Product Angle", + "RefineNodeRotateImage": "RefineNode Rotate Image", "RefineNodePreprocessMask": "RefineNode Preprocess Mask", "RefineNodeReferenceImageProcess": "RefineNode Reference Image Process", "RefineNodeRestoreMaskToOriginal": "RefineNode Restore Mask To Original",