import copy from typing import Dict, List, Tuple import torch from pytorch360convert import c2e, e2c, e2e, e2p class C2ENode: """ Cubemap To Equirectangular Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "e_img": ("IMAGE", {"default": None}), "h": ("INT", {"default": -1}), "w": ("INT", {"default": -1}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), "cube_format": ( ["stack", "dice", "horizon", "list", "dict"], {"default": "stack"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Equirectangular Image",) FUNCTION = "c2e" CATEGORY = "pytorch360convert/equirectangular" def c2e( self, e_img: torch.Tensor, h=None, w=None, padding_mode: str = "bilinear", cube_format: str = "stack", ) -> Tuple[torch.Tensor]: assert e_img.shape[0] == 6, f"Input should have 6 faces, got {e_img.shape[0]}" h = None if h < 1 else h w = None if w < 1 else w return ( c2e( e_img.squeeze(0), h=h, w=w, cube_format=cube_format, mode=padding_mode, channels_first=False, ).unsqueeze(0), ) class E2CNode: """ Equirectangular to Cubemap Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "e_img": ("IMAGE", {"default": None}), "face_width": ("INT", {"default": -1}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), "cube_format": ( ["stack", "dice", "horizon", "list", "dict"], {"default": "stack"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Cubemap Image",) FUNCTION = "e2c" CATEGORY = "pytorch360convert/equirectangular" def e2c( self, e_img: torch.Tensor, face_width: int = -1, padding_mode: str = "bilinear", cube_format: str = "stack", ) -> Tuple[torch.Tensor]: assert e_img.shape[0] == 1, ( "Only a batch size of 1 is currently" + f"supported, got {e_img.shape[0]}" ) face_width = e_img.shape[1] // 2 if face_width < 1 else face_width output = e2c( e_img.squeeze(0), face_w=face_width, mode=padding_mode, cube_format=cube_format, channels_first=False, ) output = output.unsqueeze(0) if output.dim() == 3 else output return (output,) class E2PNode: """ Equirectangular to Perspective Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "e_img": ("IMAGE", {"default": None}), "fov_deg_h": ("FLOAT", {"default": 90.0}), "fov_deg_v": ("FLOAT", {"default": 90.0}), "h_deg": ("FLOAT", {"default": 0.0}), "v_deg": ("FLOAT", {"default": 0.0}), "out_h": ("INT", {"default": 512}), "out_w": ("INT", {"default": 512}), "in_rot_deg": ("FLOAT", {"default": 0.0}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Perspective Image",) FUNCTION = "e2p" CATEGORY = "pytorch360convert/equirectangular" def e2p( self, e_img: torch.Tensor, fov_deg_h: float = 90.0, fov_deg_v: float = 90.0, h_deg: float = 0.0, v_deg: float = 0.0, out_h: int = 512, out_w: int = 512, in_rot_deg: float = 0.0, padding_mode: str = "bilinear", ) -> Tuple[torch.Tensor]: assert e_img.dim() == 4, f"e_img should have 4 dimensions, got {e_img.dim()}" return ( e2p( e_img, fov_deg=(fov_deg_h, fov_deg_v), h_deg=h_deg, v_deg=v_deg, out_hw=(out_h, out_w), in_rot_deg=in_rot_deg, mode=padding_mode, channels_first=False, ), ) class E2ENode: """ Equirectangular Rotation Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "e_img": ("IMAGE", {"default": None}), "roll": ("FLOAT", {"default": 0.0}), "h_deg": ("FLOAT", {"default": 0.0}), "v_deg": ("FLOAT", {"default": 0.0}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Rotated Image",) FUNCTION = "e2e" CATEGORY = "pytorch360convert/equirectangular" def e2e( self, e_img: torch.Tensor, roll: float = 0.0, h_deg: float = 0.0, v_deg: float = 0.0, padding_mode: str = "bilinear", ) -> Tuple[torch.Tensor]: assert e_img.dim() == 4, f"e_img should have 4 dimensions, got {e_img.dim()}" return ( e2e( e_img=e_img, h_deg=h_deg, v_deg=v_deg, roll=roll, mode=padding_mode, channels_first=False, ), ) class SplitFacesNode: """ Split a stack of cube faces for easy manipulation. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "face_stack": ("IMAGE", {"default": None}), }, } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("Front", "Right", "Back", "Left", "Up", "Down") FUNCTION = "split_faces" CATEGORY = "pytorch360convert/miscellaneous" def split_faces(self, face_stack: torch.Tensor) -> Tuple[torch.Tensor, ...]: assert ( face_stack.dim() == 4 ), f"face_stack should have 4 dimensions, got {face_stack.dim()}" assert ( face_stack.shape[0] == 6 ), f"face_stack should have 6 faces, got {face_stack.shape[0]}" outputs = [face_stack[i].unsqueeze(0) for i in range(face_stack.shape[0])] return tuple(outputs) class StackFacesNode: """ Combine multiple cube faces into a stack. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "Front": ("IMAGE", {"default": None}), "Right": ("IMAGE", {"default": None}), "Back": ("IMAGE", {"default": None}), "Left": ("IMAGE", {"default": None}), "Up": ("IMAGE", {"default": None}), "Down": ("IMAGE", {"default": None}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Cubemap stack",) FUNCTION = "stack_faces" CATEGORY = "pytorch360convert/miscellaneous" def stack_faces( self, Front: torch.Tensor, Right: torch.Tensor, Back: torch.Tensor, Left: torch.Tensor, Up: torch.Tensor, Down: torch.Tensor, ) -> Tuple[torch.Tensor]: assert ( Front.dim() == 4 and Right.dim() == 4 and Back.dim() == 4 and Left.dim() == 4 and Up.dim() == 4 and Down.dim() == 4 ) return (torch.cat([Front, Right, Back, Left, Up, Down], 0),) class RollImageNode: """ Roll an image to make face seams easier to remove or access. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "roll_x": ("INT", {"default": 0}), "roll_y": ("INT", {"default": 0}), "roll_x_by_50_percent": ( "BOOLEAN", { "default": False, "tooltip": "Ignores roll_x and roll_y. Shifts image horizontally by 50%.", }, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Rolled Image",) FUNCTION = "roll_image" CATEGORY = "pytorch360convert/miscellaneous" def roll_image( self, image: torch.Tensor, roll_x: int = 0, roll_y: int = 0, roll_x_by_50_percent: bool = False, ) -> Tuple[torch.Tensor]: assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}" if roll_x_by_50_percent: _, H, W, _ = image.shape px_half = W // 2 roll_y = 0 roll_x = px_half return (torch.roll(image, shifts=(roll_y, roll_x), dims=(1, 2)),) class C2EMaskedDiffNode: """ Cubemap To Equirectangular with Masked Diff Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "original_faces": ("IMAGE", {"default": None}), "modified_faces": ("IMAGE", {"default": None}), "original_equi": ("IMAGE", {"default": None}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), "cube_format": ( ["stack", "dice", "horizon", "list", "dict"], {"default": "stack"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Equirectangular Image",) FUNCTION = "c2e_masked_diff" CATEGORY = "pytorch360convert/mask" def c2e_masked_diff( self, original_faces: torch.Tensor, modified_faces: torch.Tensor, original_equi: torch.Tensor, padding_mode: str = "bilinear", cube_format: str = "stack", ) -> Tuple[torch.Tensor]: assert original_equi.shape[0] == 1, ( "Only a batch size of 1 is currently supported" + f"for original_equi, got {original_equi.shape[0]}" ) assert ( original_faces.shape[0] == 6 ), f"original_faces should have 6 faces, got {original_faces.shape[0]}" assert ( modified_faces.shape[0] == 6 ), f"modified_faces should have 6 faces, got {modified_faces.shape[0]}" new_equi = c2e(modified_faces, cube_format=cube_format, channels_first=False) faces_mask = original_faces != modified_faces mask_equi = c2e( faces_mask.float(), cube_format=cube_format, channels_first=False ).bool() return (torch.where(mask_equi, new_equi, original_equi),) class CropImageWithCoordsNode: """ Crop an section of an image for manipulation """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "crop_h": ("INT", {"default": 0}), "crop_w": ("INT", {"default": 0}), "crop_h2": ("INT", {"default": -1}), "crop_w2": ("INT", {"default": -1}), }, } RETURN_TYPES = ( "IMAGE", "LIST", ) RETURN_NAMES = ( "Cropped Image", "Coords", ) FUNCTION = "crop_image" CATEGORY = "pytorch360convert/miscellaneous" def crop_image( self, image: torch.Tensor, crop_h: int = 0, crop_w: int = 0, crop_h2: int = -1, crop_w2: int = -1, ) -> Tuple[torch.Tensor, List[int]]: assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}" _, H, W, _ = image.shape # Calculate the center crop indices if crop_h2 < 1: start_h = (H - crop_h) // 2 end_h = start_h + crop_h else: start_h = crop_h end_h = crop_h2 if crop_w2 < 1: start_w = (W - crop_w) // 2 end_w = start_w + crop_w else: start_w = crop_w end_w = crop_w2 cropped_tensor = image[:, start_h:end_h, start_w:end_w, :] coords = [start_h, end_h, start_w, end_w] return ( cropped_tensor, coords, ) class PasteImageWithCoordsNode: """ Add a cropped section back to the original image. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "full_image": ("IMAGE", {"default": None}), "cropped_image": ("IMAGE", {"default": None}), "coords": ("LIST", {"default": None}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Pasted Image",) FUNCTION = "paste_image" CATEGORY = "pytorch360convert/miscellaneous" def paste_image( self, full_image: torch.Tensor, cropped_image: torch.Tensor, coords: List[int] ) -> Tuple[torch.Tensor]: assert ( full_image.dim() == 4 ), f"full_image should have 4 dimensions, got {full_image.dim()}" assert ( cropped_image.dim() == 4 ), f"cropped_image should have 4 dimensions, got {cropped_image.dim()}" assert len(coords) == 4 start_h, end_h, start_w, end_w = coords full_image[:, start_h:end_h, start_w:end_w, :] = cropped_image return (full_image,) class Pad180To360Node: """ Apply padding to a 180 degree equirectangular image, to make it 360 degrees. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "fill_value": ("FLOAT", {"default": 0.0}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Padded 360 Image",) FUNCTION = "pad_180_to_360_image" CATEGORY = "pytorch360convert/equirectangular" def pad_180_to_360_image( self, image: torch.Tensor, fill_value: float = 0.0 ) -> Tuple[torch.Tensor]: assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}" image = image.permute(0, 3, 1, 2) H, W = image.shape[2:] pad_left = W // 2 pad_right = W - pad_left image_padded = torch.nn.functional.pad( image, (pad_left, pad_right), mode="constant", value=fill_value ) image_padded = image_padded.permute(0, 2, 3, 1) return (image_padded,) class Crop360To180Node: """ Crop a 360 degree equirectangular image to the central 180 degree part. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Cropped 180 Image",) FUNCTION = "crop_360_to_180_image" CATEGORY = "pytorch360convert/equirectangular" def crop_360_to_180_image(self, image: torch.Tensor) -> Tuple[torch.Tensor]: """ Crop a 360-degree equirectangular image to the central 180-degree part. Args: image (torch.Tensor): The 360-degree equirectangular image. Shape should be: (B, H, W, C). Returns: torch.Tensor: The cropped 180-degree equirectangular image. Shape will be: (B, H, W//2, C) where the width is halved. Raises: ValueError: If the input image width is less than or equal to 1. """ assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}" _, _, width, _ = image.shape # Crop the central 180-degree part by slicing the image. crop_start = width // 4 crop_end = 3 * width // 4 # Crop along the width dimension (left and right 180 degrees) cropped_img = image[:, :, crop_start:crop_end, :] return (cropped_img,) class StereoToMonoScopicNode: """ Split a stereo image into 2 monoscopic images. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "split_direction": ( ["horizontal", "vertical"], {"default": "horizontal"}, ), "larger_side": (["first", "second"], {"default": "first"}), }, } RETURN_TYPES = ( "IMAGE", "IMAGE", ) RETURN_NAMES = ( "First Image", "Second Image", ) FUNCTION = "split_stereo_image" CATEGORY = "pytorch360convert/stereo" def split_stereo_image( self, image: torch.Tensor, split_direction: str = "horizontal", larger_side: str = "first", ) -> Tuple[torch.Tensor, torch.Tensor]: """ Splits a batch of tensor images (BHWC format) into two halves either vertically or horizontally, splitting exactly down the middle. In case of odd dimensions, one side gets an extra pixel. Args: image (torch.Tensor): The tensor image batch to split. Shape should be: (B, H, W, C). split_direction (str, optional): The direction to split the image. Either "horizontal" or "vertical". Default: "horizontal" larger_side (str, optional): Specifies which side gets the extra pixel when dimensions are odd. "first" gives the extra pixel to the first side (top/left), "second" gives it to the second side (bottom/right). Default: "first" Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing the first and second halves of the split images. The shape of each tensor is: - For horizontal split: (B, H/2, W, C) - For vertical split: (B, H, W/2, C) Raises: ValueError: If `split_direction` is not "horizontal" or "vertical". """ _, height, width, channels = image.shape if split_direction == "vertical": # Split horizontally (by height) mid_point = height // 2 if height % 2 != 0: # If height is odd if larger_side == "first": first_half = image[:, : mid_point + 1, :, :] second_half = image[:, mid_point + 1 :, :, :] else: # "second" first_half = image[:, :mid_point, :, :] second_half = image[:, mid_point:, :] else: # Even height, just split in the middle first_half = image[:, :mid_point, :, :] second_half = image[:, mid_point:, :] elif split_direction == "horizontal": # Split vertically (by width) mid_point = width // 2 if width % 2 != 0: # If width is odd if larger_side == "first": first_half = image[:, :, : mid_point + 1, :] second_half = image[:, :, mid_point + 1 :, :] else: # "second" first_half = image[:, :, :mid_point, :] second_half = image[:, :, mid_point:, :] else: # Even width, just split in the middle first_half = image[:, :, :mid_point, :] second_half = image[:, :, mid_point:, :] else: raise ValueError( "Invalid split direction. Please choose" + " 'horizontal' or 'vertical'. " + f"Got {split_direction}" ) return (first_half, second_half) class MonoScopicToStereoNode: """ Merge two monoscopic images into a stereo image. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "first_image": ("IMAGE", {"default": None}), "second_image": ("IMAGE", {"default": None}), "merge_direction": ( ["horizontal", "vertical"], {"default": "horizontal"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Stereo Image",) FUNCTION = "merge_monoscopic_to_stereo" CATEGORY = "pytorch360convert/stereo" def merge_monoscopic_to_stereo( self, first_image: torch.Tensor, second_image: torch.Tensor, merge_direction: str = "horizontal", ) -> Tuple[torch.Tensor]: """ Merges two monoscopic images into a single stereo image by concatenating the two images either horizontally or vertically. Args: first_image (torch.Tensor): The left monoscopic image. Shape should be: (B, H, W, C). second_image (torch.Tensor): The right monoscopic image. Shape should be: (B, H, W, C). merge_direction (str, optional): The direction to merge the images. Either "horizontal" or "vertical". Default: "horizontal" Returns: torch.Tensor: The merged stereo image. Raises: ValueError: If `merge_direction` is not "horizontal" or "vertical". """ f_batch_size, f_height, f_width, f_channels = first_image.shape s_batch_size, s_height, s_width, s_channels = second_image.shape if merge_direction == "horizontal": assert f_batch_size == s_batch_size, ( "Batch size must match: " + f"{f_batch_size} vs {s_batch_size}" ) assert f_height == s_height, ( "Height must match: " + f"{f_height} vs {s_height}" ) assert f_channels == s_channels, ( "Channels must match: " + f"{f_channels} vs {s_channels}" ) # Concatenate horizontally (by width) stereo_image = torch.cat((first_image, second_image), dim=2) elif merge_direction == "vertical": assert f_batch_size == s_batch_size, ( "Batch size must match: " + f"{f_batch_size} vs {s_batch_size}" ) assert f_width == s_width, "Width must match: " + f"{f_width} vs {s_width}" assert f_channels == s_channels, ( "Channels must match: " + f"{f_channels} vs {s_channels}" ) # Concatenate vertically (by height) stereo_image = torch.cat((first_image, second_image), dim=1) else: raise ValueError( "Invalid split direction. Please choose" + " 'horizontal' or 'vertical'. " + f"Got {merge_direction}" ) return (stereo_image,) def _conv_forward( self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor ) -> torch.Tensor: x = torch.nn.functional.pad(x, self.padding_values_x, mode="circular") x = torch.nn.functional.pad(x, self.padding_values_y, mode="constant") return torch.nn.functional.conv2d( x, weight, bias, self.stride, (0, 0), self.dilation, self.groups ) def _apply_circular_conv2d_padding( model: torch.nn.Module, is_vae: bool = False, x_axis_only: bool = True ) -> torch.nn.Module: for layer in model.first_stage_model.modules() if is_vae else model.modules(): if isinstance(layer, torch.nn.Conv2d): if x_axis_only: layer.padding_values_x = ( layer._reversed_padding_repeated_twice[0], layer._reversed_padding_repeated_twice[1], 0, 0, ) layer.padding_values_y = ( 0, 0, layer._reversed_padding_repeated_twice[2], layer._reversed_padding_repeated_twice[3], ) layer._conv_forward = _conv_forward.__get__(layer, torch.nn.Conv2d) else: layer.padding_mode = "circular" return model class ApplyCircularConvPaddingModel: @classmethod def INPUT_TYPES(s): return { "required": { "model": ( "MODEL", {"tooltip": "Model to add circular x-axis conv2d padding to."}, ), "inplace": ( "BOOLEAN", { "default": True, "tooltip": "Modify " + "the already loaded model (True) or a copy of the model (False). " + "If True, model will have to be reloaded to restore padding to " + "the original values. Modifying inplace will use less memory.", }, ), "x_axis_only": ( "BOOLEAN", { "default": True, "tooltip": "Apply" + " circular padding only to the x-axis or to both the x and y axes.", }, ), }, } CATEGORY = "pytorch360convert/models" RETURN_TYPES = ("MODEL",) FUNCTION = "run" def run( self, model: torch.nn.Module, inplace: bool = True, x_axis_only: bool = True ) -> Tuple[torch.nn.Module]: if inplace: use_model = model else: use_model = copy.deepcopy(model) _apply_circular_conv2d_padding(use_model.model, False, x_axis_only) return (use_model,) class ApplyCircularConvPaddingVAE: @classmethod def INPUT_TYPES(cls): return { "required": { "vae": ( "VAE", {"tooltip": "VAE to add circular x-axis conv2d padding to."}, ), "inplace": ( "BOOLEAN", { "default": True, "tooltip": "Modify " + "the already loaded VAE (True) or a copy of the VAE (False). " + "If True, VAE will have to be reloaded to restore padding to " + "the original values. Modifying inplace will use less memory.", }, ), "x_axis_only": ( "BOOLEAN", { "default": True, "tooltip": "Apply" + " circular padding only to the x-axis or to both the x and y axes.", }, ), } } RETURN_TYPES = ("VAE",) FUNCTION = "run" CATEGORY = "pytorch360convert/models" def run( self, vae: torch.nn.Module, inplace: bool = True, x_axis_only: bool = True ) -> Tuple[torch.nn.Module]: if inplace: use_vae = vae else: use_vae = copy.deepcopy(vae) _apply_circular_conv2d_padding(use_vae, True, x_axis_only) return (use_vae,) def _create_center_seam_mask( x: torch.Tensor, frac_width: float = 0.10, pixel_width: int = 0, feather: int = 0 ) -> torch.Tensor: """ For a ComfyUI-style mask: shape [B, H, W], values 0 or 1, with optional feathering. Args: x (torch.Tensor): input tensor with shape [B, H, W, C]. frac_width (float, optional): fraction of input width for the vertical strip. pixel_width (int, optional): width of the seam in pixels (if > 0). feather (int, optional): pixel size of feathering on both sides of the mask. Returns: mask: torch.Tensor of shape [B, H, W] with float values 0.0 to 1.0. """ # Extract batch, height, and width from x B, H, W, *_ = x.shape # Determine strip width if pixel_width > 0: strip = pixel_width else: strip = max(1, int(W * frac_width)) x0 = (W - strip) // 2 x1 = x0 + strip # Create the mask with zeros mask = torch.zeros((B, H, W), dtype=x.dtype, device=x.device) if feather <= 0: mask[:, :, x0:x1] = 1.0 else: # Create feathered mask # Left feather region left_feather_start = max(0, x0 - feather) left_feather_end = x0 if left_feather_end > left_feather_start: feather_steps = torch.linspace( 0.0, 1.0, left_feather_end - left_feather_start, dtype=x.dtype, device=x.device, ) mask[:, :, left_feather_start:left_feather_end] = feather_steps[ None, None, : ] # Center region (full mask) mask[:, :, x0:x1] = 1.0 # Right feather region right_feather_start = x1 right_feather_end = min(W, x1 + feather) if right_feather_end > right_feather_start: feather_steps = torch.linspace( 1.0, 0.0, right_feather_end - right_feather_start, dtype=x.dtype, device=x.device, ) mask[:, :, right_feather_start:right_feather_end] = feather_steps[ None, None, : ] return mask class CreateSeamMask: """ Create a seam mask for inpainting on equirectangular images. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "frac_width": ("FLOAT", {"default": 0.10}), "pixel_width": ( "INT", { "default": 0, "tooltip": "Width of the seam in pixels." + " If > 0, frac_width is ignored in favour of width_pixel.", }, ), "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.", }, ), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Seam Mask",) FUNCTION = "run" CATEGORY = "pytorch360convert/mask" def run( self, image: torch.Tensor, frac_width: float = 0.10, pixel_width: int = 0, 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 seam_mask = _create_center_seam_mask( image, frac_width=frac_width, pixel_width=pixel_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,) class E2Face: """ Equirectangular to Face """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "e_img": ("IMAGE", {"default": None}), "face_width": ("INT", {"default": -1}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), "cube_face": ( ["Up", "Down", "Right", "Left", "Front", "Back"], {"default": "Front"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Face Image",) FUNCTION = "run_e2face" CATEGORY = "pytorch360convert/equirectangular" def run_e2face( self, e_img: torch.Tensor, face_width: int = -1, padding_mode: str = "bilinear", cube_face: str = "Front", ) -> Tuple[torch.Tensor]: B, H, W, C = e_img.shape outputs = [] for i in range(B): single_img = e_img[i] # [H,W,C] # Determine face width face_w = H // 2 if face_width < 1 else face_width # Convert single equirectangular image to cubemap dict cubemap = e2c( single_img, face_w=face_w, mode=padding_mode, cube_format="dict", channels_first=False, ) # Pick requested face face_tensor = cubemap[cube_face] # [face_w, face_w, C] outputs.append(face_tensor.unsqueeze(0)) # [1,H,W,C] # Concatenate into [B,H,W,C] output_batch = torch.cat(outputs, dim=0) return (output_batch,) def _create_centered_circle_mask( x: torch.Tensor, circle_radius: float, feather: int = 0, pixel_radius: int = 0, ) -> torch.Tensor: """ Create a centered circle mask with optional feathering. Args: x (torch.Tensor): Reference tensor with shape (B, C, H, W). The mask will copy device, dtype, and batch size from x. circle_radius (float): Fraction (0.0–1.0) of max possible radius (min(H,W)/2). 1.0 = circle touches edges. feather (int, optional): Feather width in pixels, extending outward from circle_radius. 0 = hard edge, >0 = smooth falloff. Default: 0 pixel_radius (int, optional): Explicit radius in pixels. If >0, overrides circle_radius. Default: 0 Returns: torch.Tensor: Mask tensor of shape (1, C, H, W), values in [0.0, 1.0]. """ _, C, H, W = x.shape size = min(H, W) # Circle radius in pixels max_radius = size / 2.0 if pixel_radius > 0: inner_radius = float(pixel_radius) else: inner_radius = circle_radius * max_radius outer_radius = inner_radius + feather # Coordinate grid yy, xx = torch.meshgrid( torch.arange(size, device=x.device), torch.arange(size, device=x.device), indexing="ij", ) center = size // 2 dist = torch.sqrt((xx - center) ** 2 + (yy - center) ** 2) # Base mask mask = torch.zeros_like(dist, dtype=x.dtype, device=x.device) # Inside circle = 1.0 mask[dist <= inner_radius] = 1.0 # Feather transition if feather > 0: transition_zone = (dist > inner_radius) & (dist <= outer_radius) mask[transition_zone] = 1.0 - (dist[transition_zone] - inner_radius) / feather # Expand to BCHW and center in full HxW mask_full = torch.zeros((1, C, H, W), dtype=x.dtype, device=x.device) y0 = (H - size) // 2 x0 = (W - size) // 2 mask_full[:, :, y0 : y0 + size, x0 : x0 + size] = mask return mask_full class CreatePoleMask: """ Create a pole mask for inpainting on equirectangular images. """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "circle_radius": ("FLOAT", {"default": 0.10}), "pixel_radius": ( "INT", { "default": 0, "tooltip": "Size of the radius in pixels." + " If > 0, circle_radius is ignored in favour of pixel_radius.", }, ), "feather": ("INT", {"default": 0}), "mode": (["face", "equirectangular"], {"default": "face"}), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Pole Mask",) FUNCTION = "run" CATEGORY = "pytorch360convert/mask" def run( self, image: torch.Tensor, circle_radius: float = 0.10, pixel_radius: int = 0, feather: int = 0, mode: str = "face", ) -> Tuple[torch.Tensor]: assert image.dim() == 4, f"Image should have 4 dimensions, got {image.shape}" if mode == "face": image = image.permute(0, 3, 1, 2) # BHWC -> BCHW output_mask = _create_centered_circle_mask( image, circle_radius=circle_radius, pixel_radius=pixel_radius, feather=feather, ) output_mask = output_mask else: output_mask = [] for im in image: # im: [H,W,C] # Convert single equirectangular image to cubemap dict cubemap = e2c( torch.zeros_like(im), cube_format="dict", channels_first=False ) # Apply circle mask on Up and Down faces for face in ["Up", "Down"]: face_tensor = cubemap[face] # [H_face, W_face, C] face_tensor = face_tensor.unsqueeze(0).permute(0, 3, 1, 2) face_tensor = _create_centered_circle_mask( face_tensor, circle_radius=circle_radius, pixel_radius=pixel_radius, feather=feather, ) cubemap[face] = face_tensor.permute(0, 2, 3, 1).squeeze(0) # Convert back to equirectangular new_equi = c2e( cubemap, cube_format="dict", channels_first=False ).unsqueeze( 0 ) # add batch dim output_mask.append(new_equi) # Stack all results into a single tensor [B,H,W,C] output_mask = torch.cat(output_mask).permute(0, 3, 1, 2) return (output_mask[:, 0, ...],) class Face2E: """ Face To Equirectangular Node """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "face": ( ["Up", "Down", "Front", "Right", "Left", "Back"], {"default": "Down"}, ), "base_equi_color": ("FLOAT", {"default": 0.0}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Equirectangular Image",) FUNCTION = "run_face2e" CATEGORY = "pytorch360convert/equirectangular" def run_face2e( self, image: torch.Tensor, face: str = "Down", base_equi_color: float = 0.0, padding_mode: str = "bilinear", ) -> Tuple[torch.Tensor]: assert image.dim() == 4, f"Image should have 4 dimensions, got {image.shape}" output_image = [] for f_image in image: cubemap_dict = {} for face_name in ["Front", "Right", "Back", "Left", "Up", "Down"]: if face_name != face: cubemap_dict[face_name] = torch.ones_like(f_image) * base_equi_color else: cubemap_dict[face_name] = f_image output_image += [ c2e( cubemap=cubemap_dict, cube_format="dict", mode=padding_mode, channels_first=False, ) ] 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/mask" 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)),) class FaceMask2E: """ Face Mask To Equirectangular """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "mask": ("MASK", {"default": None}), "face": ( ["Up", "Down", "Front", "Right", "Left", "Back"], {"default": "Down"}, ), "base_equi_color": ("FLOAT", {"default": 0.0}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Equirectangular Mask",) FUNCTION = "run_facemask2e" CATEGORY = "pytorch360convert/mask" def run_facemask2e( self, mask: torch.Tensor, face: str = "Down", base_equi_color: float = 0.0, padding_mode: str = "bilinear", ) -> Tuple[torch.Tensor]: assert mask.dim() == 3, f"Image should have 3 dimensions, got {mask.shape}" output_mask = [] for f_mask in mask: f_mask = f_mask = f_mask[None, ...] cubemap_dict = {} for face_name in ["Front", "Right", "Back", "Left", "Up", "Down"]: if face_name != face: cubemap_dict[face_name] = torch.ones_like(f_mask) * base_equi_color else: cubemap_dict[face_name] = mask.reshape(1, *mask.shape[1:]) output_mask += [ c2e( cubemap=cubemap_dict, cube_format="dict", mode=padding_mode, channels_first=True, ) ] return (torch.cat(output_mask),) class EMask2Face: """ Equirectangular Mask to Face """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "mask": ("MASK", {"default": None}), "face_width": ("INT", {"default": -1}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), "cube_face": ( ["Up", "Down", "Right", "Left", "Front", "Back"], {"default": "Front"}, ), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Face MASK",) FUNCTION = "run_emask2face" CATEGORY = "pytorch360convert/mask" def run_emask2face( self, mask: torch.Tensor, face_width: int = -1, padding_mode: str = "bilinear", cube_face: str = "Front", ) -> Tuple[torch.Tensor]: B, H, W = mask.shape outputs = [] for i in range(B): singlmask = mask[i : i + 1][..., None] # [H,W,C] # Determine face width face_w = H // 2 if face_width < 1 else face_width # Convert single equirectangular mask to cubemap dict cubemap = e2c( singlmask, face_w=face_w, mode=padding_mode, cube_format="dict", channels_first=False, ) # Pick requested face face_tensor = cubemap[cube_face] # [face_w, face_w, C] outputs.append(face_tensor) # [1,H,W,C] # Concatenate into [B,H,W,C] output_batch = torch.cat(outputs, dim=0)[..., 0] return (output_batch,) class MaskE2ENode: """ Mask Equirectangular Rotation """ @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "mask": ("MASK", {"default": None}), "roll": ("FLOAT", {"default": 0.0}), "h_deg": ("FLOAT", {"default": 0.0}), "v_deg": ("FLOAT", {"default": 0.0}), "padding_mode": ( ["bilinear", "bicubic", "nearest"], {"default": "bilinear"}, ), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Rotated Mask",) FUNCTION = "mask_e2e" CATEGORY = "pytorch360convert/mask" def mask_e2e( self, mask: torch.Tensor, roll: float = 0.0, h_deg: float = 0.0, v_deg: float = 0.0, padding_mode: str = "bilinear", ) -> Tuple[torch.Tensor]: assert mask.dim() == 3, f"mask should have 3 dimensions, got {mask.dim()}" return ( e2e( e_img=mask[..., None], h_deg=h_deg, v_deg=v_deg, roll=roll, mode=padding_mode, channels_first=False, )[..., 0], ) class Create180To360Mask: @classmethod def INPUT_TYPES(s) -> Dict: return { "required": { "image": ("IMAGE", {"default": None}), "input_mode": ( ["360", "180"], {"default": "180"}, ), "feather": ("INT", {"default": 0}), }, } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("Equirectangular Mask",) FUNCTION = "mask_180_to_360" CATEGORY = "pytorch360convert/mask" def mask_180_to_360( self, image: torch.Tensor, input_mode: str = "180", feather: int = 0 ) -> Tuple[torch.Tensor]: assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}" _, H, W, _ = image.shape # For 360 input, the valid region is only half width if input_mode == "360": W = W // 2 pad_left = W // 2 pad_right = W - pad_left total_width = W + pad_left + pad_right # == 2*W # Start with zeros everywhere mask = torch.zeros(H, total_width, dtype=image.dtype, device=image.device) # Fill the main region with 1.0 mask[:, pad_left : pad_left + W] = 1.0 if feather > 0: ramp = torch.linspace( 0, 1, steps=feather + 1, device=image.device, dtype=image.dtype )[1:] # Left feather (in the padded zero region, approaching the mask) mask[:, pad_left - feather : pad_left] = ramp # Right feather (in the padded zero region, approaching the mask) mask[:, pad_left + W : pad_left + W + feather] = ramp.flip(0) # [1, H, W] mask tensor return (mask.unsqueeze(0),)