diff --git a/README.md b/README.md index 0ae4208..d3bd0a6 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # ComfyUI-Environment-Visualizer -This ComfyUI node pack allows the user to take a panoramic image and a corresponding depth map, and turn them into a 3D environment, which they can view in an immersive WebXR environment. +This ComfyUI node pack allows the user to take a panoramic image and a corresponding depth map, and turn them into a 3D environment, which they can view in an immersive WebXR environment. It also includes utility nodes to assist with creating equirectangular images. ## Operational overview When this node pack is installed, ComfyUI will launch a web server on the same machine when ComfyUI is initialized. By default, the server will be hosted on port 4443, but this can be changed by modifying the value of `SERVER_PORT` at the top of `server.py`. As WebXR requires an HTTPS environment to run, a self-signed certificate is generated on first launch, which will display a warning when loaded on most browsers. The web page includes a dropdown to pick from previously saved environments, as well as several controls to adjust the presentation of the environment. The user can turn left or right in the environment with the joysticks on a VR controller, or with A/D and Left/Right on a keyboard. As ComfyUI will always launch this web server on start whenever the node pack is active, you may consider disabling the pack when it is not in use. @@ -7,9 +7,10 @@ When this node pack is installed, ComfyUI will launch a web server on the same m ### Included ComfyUI nodes - **Environment Visualizer**: This output node accepts two image inputs, a 2:1 image for the panoramic skybox, and a 2:1 image for the depth map. The node includes a text input for the environment name, and a toggle to automatically launch the generated visualization in a new tab when the workflow has completed. Running this node will automatically save the texture and depth map images into a new subdirectory in the `/environments` directory with the provided name. The `depth` input is optional, and the environment will be displayed as a flat skybox if it is omitted. - **Interpolate Edges**: This node takes an image, and linearly interpolates the edge pixel values along the vertical edges such that they continuously blend to the opposite edge. This is intended for smoothly blending the seams of the generated depth maps to remove discontinuities. The smoothing will be applied over the given number of pixels on each side. While this approach will work on the panoramic texture or any other image, it is more noticeable in color, and is primarily intended for use with the depth map. +- **MapEquirectangular**: This node accepts an image, alongside camera parameters such as its FOV and orientation, and projects the image onto an equirectangular image of a given width. Pixels in the resulting equirectangular image that are not mapped from the input image will be black and have zero opacity. This allows the image to be split into both the RGB texture and mask to simplify inpainting. Note that processing the projection with a yaw of 180 degrees will split the image across the horizontal edges, which may help with creating seamless panoramas. ## Generation tips -Popular image diffusion models such as Stable Diffusion and FLUX are capable of generating panoramic style images, though they may benefit greatly from a checkpoint or LoRA trained specifically on 360 degree panoramas. Additionally, including phrases in the prompt such as 'HDRI', 'panoramic image of...', or '360HDR' may improve your results. All images should be generated with a 2:1 aspect ratio, such as 2048x1024. There are many existing tools available in ComfyUI to generate depth maps from images, such as Marigold Depth Estimation, or Depth Anything. +Popular image diffusion models such as Stable Diffusion and FLUX are capable of generating panoramic style images, though they may benefit greatly from a checkpoint or LoRA trained specifically on 360 degree panoramas. Additionally, including phrases in the prompt such as 'HDRI', 'panoramic image of...', or '360HDR' may improve your results. All images should be generated with a 2:1 aspect ratio, such as 2048x1024. There are many existing tools available in ComfyUI to generate depth maps from images, such as Marigold Depth Estimation and Depth Anything. ## Standalone use This WebXR server included in this repository can be used as a standalone utility by running `server.py` manually. The dropdown on the page will display the names of all subdirectories in `/environments`, and attempt to load the `skybox.png` and `depth.png` images in the selected subdirectory. Although the skybox texture is required, the depth map is optional. While running the Environment Visualizer node in ComfyUI will automatically save these files to the correct path, any correctly named images placed in the environments directory can be loaded and viewed. \ No newline at end of file diff --git a/__init__.py b/__init__.py index 89cc691..2105e11 100644 --- a/__init__.py +++ b/__init__.py @@ -1,14 +1,16 @@ from .node import EnvironmentVisualizer from .node import InterpolateEdges +from .node import MapEquirectangular from .server import run_https_server import threading NODE_CLASS_MAPPINGS = { "EnvironmentVisualizer": EnvironmentVisualizer, - "InterpolateEdges": InterpolateEdges + "InterpolateEdges": InterpolateEdges, + "MapEquirectangular": MapEquirectangular } WEB_DIRECTORY = "./web" server_thread = threading.Thread(target=run_https_server) -server_thread.start() +server_thread.start() \ No newline at end of file diff --git a/map_equirectangular.py b/map_equirectangular.py new file mode 100644 index 0000000..205533a --- /dev/null +++ b/map_equirectangular.py @@ -0,0 +1,154 @@ +import torch +import torch.nn.functional as F +import math + + +def map_equirectangular(input_tensor, HFOV, yaw, pitch, roll, output_width=4096): + """ + Maps an input image tensor to an equirectangular panoramic image using PyTorch tensors. + + Parameters: + input_tensor (torch.Tensor): Input image as a tensor with shape [H, W, C] in RGB or RGBA format. + HFOV (float): Horizontal Field of View in degrees. + yaw (float): Yaw rotation in degrees. + pitch (float): Pitch rotation in degrees. + roll (float): Roll rotation in degrees. + output_width (int): Width of the output equirectangular image (Height will be output_width // 2). + + Returns: + equirect_image (torch.Tensor): Equirectangular image tensor with shape [output_width // 2, output_width, 4]. + """ + if input_tensor.ndim != 3: + raise ValueError("Input tensor must be a 3-dimensional array [H, W, C].") + + H_in, W_in, C = input_tensor.shape + + # Handle images with or without an alpha channel + if C == 4: + # Input has alpha channel; preserve RGB and ignore input alpha + input_image = input_tensor[:, :, :3].clone() + elif C == 3: + input_image = input_tensor.clone() + else: + raise ValueError("Input tensor must have 3 (RGB) or 4 (RGBA) channels.") + + # Calculate Vertical Field of View (VFOV) to match FOV per pixel + VFOV = HFOV * (H_in / W_in) + + # Convert FOV from degrees to radians + HFOV_rad = torch.deg2rad(torch.tensor(HFOV, dtype=input_tensor.dtype, device=input_tensor.device)) + VFOV_rad = torch.deg2rad(torch.tensor(VFOV, dtype=input_tensor.dtype, device=input_tensor.device)) + + # Compute focal lengths + fx = (W_in / 2) / torch.tan(HFOV_rad / 2) + fy = (H_in / 2) / torch.tan(VFOV_rad / 2) + + # Principal point (assuming centered) + cx = W_in / 2 + cy = H_in / 2 + + # Compute the rotation matrix from yaw, pitch, roll + def rotation_matrix(yaw, pitch, roll): + # Convert angles from degrees to radians + yaw_rad = math.radians(yaw) + pitch_rad = math.radians(pitch) + roll_rad = math.radians(roll) + + # Rotation matrices around x, y, z axes + Rx = torch.tensor([ + [1, 0, 0], + [0, math.cos(pitch_rad), -math.sin(pitch_rad)], + [0, math.sin(pitch_rad), math.cos(pitch_rad)] + ], dtype=input_tensor.dtype, device=input_tensor.device) + + Ry = torch.tensor([ + [math.cos(yaw_rad), 0, math.sin(yaw_rad)], + [0, 1, 0], + [-math.sin(yaw_rad), 0, math.cos(yaw_rad)] + ], dtype=input_tensor.dtype, device=input_tensor.device) + + Rz = torch.tensor([ + [math.cos(roll_rad), -math.sin(roll_rad), 0], + [math.sin(roll_rad), math.cos(roll_rad), 0], + [0, 0, 1] + ], dtype=input_tensor.dtype, device=input_tensor.device) + + # Combined rotation matrix + R = Rz @ Ry @ Rx + return R + + # Compute rotation matrix and ensure it's contiguous + R = rotation_matrix(yaw, pitch, roll).T # Transpose for inverse rotation + + # Equirectangular image dimensions (Height is half of Width) + W_out = output_width + H_out = output_width // 2 # Enforce 2:1 aspect ratio + + # Create a meshgrid for the equirectangular image + theta = (torch.linspace(0, W_out - 1, W_out, dtype=input_tensor.dtype, device=input_tensor.device) / W_out) * 2 * math.pi - math.pi # theta from -π to π + phi = (0.5 - (torch.linspace(0, H_out - 1, H_out, dtype=input_tensor.dtype, device=input_tensor.device) / H_out)) * math.pi # phi from -π/2 to π/2 + + # Use torch.meshgrid with proper ordering to get [H_out, W_out] + phi_grid, theta_grid = torch.meshgrid(phi, theta, indexing='ij') # Shape: [H_out, W_out] + + # Spherical to Cartesian coordinates (direction vectors) + x_s = torch.cos(phi_grid) * torch.sin(theta_grid) # Shape: [H_out, W_out] + y_s = -torch.sin(phi_grid) # Shape: [H_out, W_out] + z_s = torch.cos(phi_grid) * torch.cos(theta_grid) # Shape: [H_out, W_out] + + # Stack into direction vectors + dirs = torch.stack((x_s, y_s, z_s), dim=-1) # Shape: [H_out, W_out, 3] + + # Rotate direction vectors to camera coordinate system + dirs_cam = torch.matmul(dirs, R) # Shape: [H_out, W_out, 3] + dx_c, dy_c, dz_c = dirs_cam[..., 0], dirs_cam[..., 1], dirs_cam[..., 2] + + # Compute valid_mask before division to avoid divide by zero + epsilon = torch.tensor(1e-6, dtype=input_tensor.dtype, device=input_tensor.device) + valid_mask = dz_c > epsilon # Points in front of the camera + + # Compute x_im and y_im + x_im = (dx_c / dz_c) * fx + cx + y_im = (dy_c / dz_c) * fy + cy + + # Update valid_mask with x_im and y_im in valid image range + valid_mask &= (x_im >= 0) & (x_im < W_in) & (y_im >= 0) & (y_im < H_in) + valid_mask &= torch.isfinite(x_im) & torch.isfinite(y_im) + + # Prepare grid for grid_sample + # Normalize x_im and y_im to [-1, 1] + grid_x = (x_im / (W_in - 1)) * 2 - 1 # Shape: [H_out, W_out] + grid_y = (y_im / (H_in - 1)) * 2 - 1 # Shape: [H_out, W_out] + + # Ensure grid_x and grid_y are within [-1, 1] + grid_x = torch.clamp(grid_x, -1.0, 1.0) + grid_y = torch.clamp(grid_y, -1.0, 1.0) + + # Stack to create grid of shape [1, H_out, W_out, 2] + grid = torch.stack((grid_x, grid_y), dim=-1) # Shape: [H_out, W_out, 2] + grid = grid.unsqueeze(0) # Shape: [1, H_out, W_out, 2] + + # Ensure grid is contiguous + grid = grid.contiguous() + + # Prepare input image tensor + # Convert to float and permute to [C, H, W] + input_image = input_image.to(dtype=input_tensor.dtype, device=input_tensor.device).permute(2, 0, 1).unsqueeze(0) # Shape: [1, C, H_in, W_in] + + # Perform remapping using grid_sample + remapped = F.grid_sample(input_image, grid, mode='bilinear', padding_mode='zeros', align_corners=True) # Shape: [1, C, H_out, W_out] + + # Create an alpha channel based on valid_mask + alpha = valid_mask.unsqueeze(0).unsqueeze(0).float() # Shape: [1, 1, H_out, W_out] + + # Mask the remapped RGB image with the alpha channel to set invalid regions to zero + remapped = remapped * alpha # Zero out invalid regions + + # Combine the remapped RGB image with the alpha channel to create RGBA image + equirect_image = torch.cat((remapped, alpha), dim=1) # Shape: [1, C, H_out, W_out] + + # Squeeze the batch dimension and permute to [H_out, W_out, C+1] + equirect_image = equirect_image.squeeze(0).permute(1, 2, 0) # Shape: [H_out, W_out, C] + + equirect_image = equirect_image.clamp(0, 1) + return equirect_image \ No newline at end of file diff --git a/node.py b/node.py index ffbdc5b..a527123 100644 --- a/node.py +++ b/node.py @@ -1,12 +1,15 @@ -from .server import get_lan_ip from .server import SERVER_PORT +from .map_equirectangular import map_equirectangular + from server import PromptServer + from aiohttp import web from PIL import Image import os import numpy as np import time import re +import torch @PromptServer.instance.routes.post("/get_url") @@ -14,6 +17,50 @@ async def get_url(_): return web.json_response({"port": str(SERVER_PORT)}) +class MapEquirectangular: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "equirectangular_width": ("INT", {"default": 2048}), + "hfov": ("FLOAT", {"default": 60.0, "min": 0.0, "max": 180.0, "step": 1.0}), + "yaw": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0}), + "pitch": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 1.0}), + "roll": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0}) + }, + } + + RETURN_TYPES = ("IMAGE", ) + OUTPUT_NODE = False + FUNCTION = "map" + CATEGORY = "image/equirectangular" + DESCRIPTION = "Takes an image and some camera parameters, and projects it onto an equirectangular image." + + def map(self, image, equirectangular_width, hfov, yaw, pitch, roll): + + B = image.shape[0] + processed_images = [] + + for i in range(B): + # Process the image using the method + processed_image = map_equirectangular( + image[i], + hfov, + yaw, + pitch, + roll, + equirectangular_width + ) + # Append the processed image to the list + processed_images.append(processed_image) + + # Aggregate the processed images into a tensor of shape [B, H, W, C] + output_tensor = torch.stack(processed_images, dim=0) + + return (output_tensor,) + class EnvironmentVisualizer: @classmethod @@ -32,7 +79,7 @@ class EnvironmentVisualizer: RETURN_TYPES = () OUTPUT_NODE = True FUNCTION = "save_environment" - CATEGORY = "image" + CATEGORY = "image/equirectangular" DESCRIPTION = "Saves the texture and depth map, to be viewed in an immersive WebXR environment." save_directory = os.path.join(os.path.dirname(__file__), 'environments') @@ -83,8 +130,7 @@ class EnvironmentVisualizer: } return { "ui": completion_data } else: - return {} - + return {} class InterpolateEdges: @@ -100,7 +146,7 @@ class InterpolateEdges: RETURN_TYPES = ("IMAGE",) FUNCTION = "interpolate_edges" - CATEGORY = "image" + CATEGORY = "image/equirectangular" DESCRIPTION = "Make the vertical edges of the given images blend seamlessly, using linear interpolation. Works best with depth maps." def interpolate_edges(self, image, distance): diff --git a/pyproject.toml b/pyproject.toml index b1a12fc..8e28228 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-environment-visualizer" description = "This ComfyUI node pack allows the user to take a panoramic image and a corresponding depth map, and turn them into a 3D environment, which they can view in an immersive WebXR environment." -version = "1.0.0" +version = "1.1.0" license = {file = "LICENSE"} [project.urls] diff --git a/server.py b/server.py index cf9136a..9f1768f 100644 --- a/server.py +++ b/server.py @@ -57,7 +57,7 @@ def run_https_server(): # Wrap the socket httpd.socket = context.wrap_socket(httpd.socket, server_side=True) - print(f"Serving HTTPS on https://{get_lan_ip()}:{SERVER_PORT}") + print(f"Serving Environment Visualizer HTTPS on https://{get_lan_ip()}:{SERVER_PORT}") httpd.serve_forever() def get_lan_ip():