From 2fb5b784a854cb4368b17e48c70b4e5575be3675 Mon Sep 17 00:00:00 2001 From: ProGamerGov Date: Thu, 21 Aug 2025 13:32:25 -0600 Subject: [PATCH] Add Roll Mask node --- ReadMe.md | 6 +++++ __init__.py | 3 +++ nodes.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++--- pyproject.toml | 2 +- 4 files changed, 75 insertions(+), 4 deletions(-) diff --git a/ReadMe.md b/ReadMe.md index fd0f5bc..3bf6f47 100644 --- a/ReadMe.md +++ b/ReadMe.md @@ -94,6 +94,12 @@ The Roll Image node rolls an image along the height dimension (y axis) or the wi If you set 'roll_x_by_50_percent' to True, then any x and y axis values are ignored, and the image is shifted horizontally to be 50%. This is equivalent to a 180 degree rotation on an equirectangular image. +### Roll Mask Axes + +The Roll Mask node rolls an mask along the height dimension (y axis) or the width dimension (x axis). This can help make artifacts like seams more accessible. Note that y axis rotations are not correct for final equirectangular outputs. Use the E2E node instead for y axis rotations. + +If you set 'roll_x_by_50_percent' to True, then any x and y axis values are ignored, and the image is shifted horizontally to be 50%. This is equivalent to a 180 degree rotation on an equirectangular image. + ### Crop Image with Coords & Paste Image with Coords These two nodes allow you to crop a section of an image out to make things like in-painting faster and more memory efficient, before pasting the cropped section back into the original image. diff --git a/__init__.py b/__init__.py index 4fbee7a..a82b51e 100644 --- a/__init__.py +++ b/__init__.py @@ -16,6 +16,7 @@ from .nodes import ( Pad180To360Node, PasteImageWithCoordsNode, RollImageNode, + RollMaskNode, SplitFacesNode, StackFacesNode, StereoToMonoScopicNode, @@ -42,6 +43,7 @@ NODE_CLASS_MAPPINGS = { "Create Pole Mask": CreatePoleMask, "Equirectangular to Face": E2Face, "Face to Equirectangular": Face2E, + "Roll Mask Axes": RollMaskNode, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -65,4 +67,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Create Pole Mask": "Create Pole Mask", "Equirectangular to Face": "Equirectangular to Face", "Face to Equirectangular": "Face to Equirectangular", + "Roll Mask Axes": "Roll Mask Axes", } diff --git a/nodes.py b/nodes.py index 6eddf94..feb053f 100644 --- a/nodes.py +++ b/nodes.py @@ -971,7 +971,18 @@ class CreateSeamMask: "required": { "image": ("IMAGE", {"default": None}), "seam_mask_width": ("FLOAT", {"default": 0.10}), - "feather": ("INT", {"default": 0}), + "feather": ( + "INT", + {"default": 0, "tooltip": "Pixel size of the feathering."}, + ), + "roll_x_by_50_percent": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Shifts the output mask horizontally by 50%." + + " Equivalent to a 180 degree rotation on an equirectangular image.", + }, + ), }, } @@ -987,12 +998,17 @@ class CreateSeamMask: image: torch.Tensor, seam_mask_width: float = 0.10, feather: int = 0, + roll_x_by_50_percent: bool = False, ) -> Tuple[torch.Tensor]: assert image.dim() == 4, "Image should have 4 dimensions" _, H, W, _ = image.shape px_half = W // 2 - image_rolled = torch.roll(image, shifts=(0, px_half), dims=(1, 2)) - seam_mask = _create_center_seam_mask(image, frac_width=seam_mask_width, feather=feather) + seam_mask = _create_center_seam_mask( + image, frac_width=seam_mask_width, feather=feather + ) + + if roll_x_by_50_percent: + seam_mask = torch.roll(seam_mask, shifts=(0, px_half), dims=(1, 2)) return (seam_mask,) @@ -1237,3 +1253,49 @@ class Face2E: ) ] return (torch.stack(output_image),) + + +class RollMaskNode: + """ + Roll a mask to make face seams easier to remove or access. + """ + + @classmethod + def INPUT_TYPES(s) -> Dict: + return { + "required": { + "mask": ("MASK", {"default": None}), + "roll_x": ("INT", {"default": 0}), + "roll_y": ("INT", {"default": 0}), + "roll_x_by_50_percent": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Shifts the input mask horizontally by 50%." + + " Equivalent to a 180 degree rotation on an equirectangular image.", + }, + ), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("Rolled Mask",) + + FUNCTION = "roll_mask" + + CATEGORY = "pytorch360convert" + + def roll_mask( + self, + mask: torch.Tensor, + roll_x: int = 0, + roll_y: int = 0, + roll_x_by_50_percent: bool = False, + ) -> Tuple[torch.Tensor]: + assert mask.dim() == 3, f"mask should have 3 dimensions, got {mask.dim()}" + if roll_x_by_50_percent: + _, H, W = mask.shape + px_half = W // 2 + roll_y = 0 + roll_x = px_half + return (torch.roll(mask, shifts=(roll_y, roll_x), dims=(1, 2)),) diff --git a/pyproject.toml b/pyproject.toml index 8bcad61..ad1aa6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-pytorch360convert" description = "A collection of custom nodes for working with and converting between and working with 360 degree equirectangular images, cubemap, and perspective images. Also includes Circular padding Conv2d options for models and VAEs, along with nodes to fix seams and polar artifacts. Panoramic 360 images are also sometimes known as VR photography (virtual reality), HDRI environments (ex: skyboxes), image spheres, spherical images, 360 pano." -version = "1.0.8" +version = "1.0.9" license = {file = "LICENSE"} dependencies = ["pytorch360convert"]