282 lines
10 KiB
Python
282 lines
10 KiB
Python
import base64
|
|
from io import BytesIO
|
|
from typing import Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
class PanoramaViewerNode:
|
|
"""A ComfyUI node that provides an interactive 360-degree panorama viewer for
|
|
equirectangular images.
|
|
|
|
This node takes an input image tensor and displays it in a Three.js-based
|
|
panoramic viewer that allows users to:
|
|
- Pan around the 360-degree view using mouse drag
|
|
- Zoom in/out using the mouse wheel
|
|
- View the panorama with proper equirectangular projection
|
|
|
|
The viewer automatically handles different image formats and sizes, including:
|
|
- Batch processing (takes first image from batch)
|
|
- Grayscale to RGB conversion
|
|
- Image resizing for performance optimization
|
|
- Float to uint8 conversion
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define the input parameters for the node.
|
|
|
|
Returns:
|
|
dict: Dictionary containing required input parameters:
|
|
- images (torch.Tensor): Input tensor containing the panoramic image.
|
|
- max_width (int): Maximum dimension for resizing. Default: 4096
|
|
Set to -1 for no resizing.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"max_width": (
|
|
"INT",
|
|
{
|
|
"default": 4096,
|
|
"tooltip": "The max width to use. Images larger than the"
|
|
+ " specified value will be resized. Larger sizes may run"
|
|
+ " slower. Set to -1 for no resizing.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "view_pano"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "pytorch360convert"
|
|
|
|
def __init__(self) -> None:
|
|
self.type = "pano"
|
|
|
|
def view_pano(
|
|
self, images: torch.Tensor, max_width: int = 4096
|
|
) -> Dict[str, Dict[str, str]]:
|
|
"""Process and display the panoramic image in the viewer.
|
|
|
|
This method handles the conversion of the input tensor to a viewable format by:
|
|
1. Extracting the first image if dealing with a batch
|
|
2. Converting the tensor to a numpy array
|
|
3. Converting float values to uint8 if necessary
|
|
4. Converting grayscale to RGB if necessary
|
|
5. Resizing the image if it exceeds max_width
|
|
6. Converting the processed image to a base64-encoded PNG
|
|
|
|
Args:
|
|
images (torch.Tensor): Input tensor containing the panoramic image(s).
|
|
Should be in format (B, H, W, C) or (H, W, C).
|
|
max_width (int, optional): Maximum width for resizing. Images will not
|
|
resized if they are smaller than the specified size. Set to -1 to
|
|
disable resizing. Default: 4096
|
|
|
|
Returns:
|
|
Dict[str, Dict[str, str]]: Dictionary containing the UI update information
|
|
with the base64-encoded PNG image data.
|
|
"""
|
|
# Handle batch dimension
|
|
if len(images.shape) == 4:
|
|
image = images[0]
|
|
else:
|
|
image = images
|
|
|
|
# Convert to numpy and proper format
|
|
image_np = image.cpu().numpy()
|
|
|
|
# Convert to uint8 if needed
|
|
if image_np.dtype != np.uint8:
|
|
image_np = (image_np * 255).astype(np.uint8)
|
|
|
|
# Handle grayscale images
|
|
if len(image_np.shape) == 2 or image_np.shape[2] == 1:
|
|
image_np = np.repeat(image_np[..., np.newaxis], 3, axis=2)
|
|
|
|
# Convert to PIL Image
|
|
pil_image = Image.fromarray(image_np)
|
|
|
|
# Optionally resize image
|
|
if max_width > 0 and (
|
|
pil_image.size[0] > max_width or pil_image.size[1] > max_width
|
|
):
|
|
new_size = tuple(
|
|
[int(max_width * x / max(pil_image.size)) for x in pil_image.size]
|
|
)
|
|
pil_image = pil_image.resize(new_size, resample=Image.Resampling.LANCZOS) # type: ignore[arg-type]
|
|
|
|
# Save to BytesIO
|
|
buffered = BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
|
|
# Get base64 string
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|
return {"ui": {"pano_image": f"data:image/png;base64,{img_str}"}}
|
|
|
|
|
|
class PanoramaVideoViewerNode:
|
|
"""A ComfyUI node that provides an interactive 360-degree panorama viewer for
|
|
equirectangular video content.
|
|
|
|
This node takes a video file path or a sequence of image frames and displays it in
|
|
a Three.js-based panoramic video viewer that allows users to:
|
|
- Pan around the 360-degree view using mouse drag
|
|
- Zoom in/out using the mouse wheel
|
|
- Play, pause, and scrub through the video timeline
|
|
- View the panorama video with proper equirectangular projection
|
|
|
|
The viewer automatically handles different video formats and sizes, including:
|
|
- Batch processing (takes frames from batch)
|
|
- Proper frame sequencing and timing
|
|
- Image resizing for performance optimization
|
|
- Float to uint8 conversion
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define the input parameters for the node.
|
|
|
|
Returns:
|
|
dict: Dictionary containing required input parameters:
|
|
- video_frames (torch.Tensor): Input tensor containing a batch of video frames.
|
|
- fps (int): Frames per second for the video playback.
|
|
- max_width (int): Maximum dimension for resizing. Default: 2048
|
|
Set to -1 for no resizing.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"video_frames": ("IMAGE",),
|
|
"fps": (
|
|
"INT",
|
|
{
|
|
"default": 30,
|
|
"min": 1,
|
|
"max": 120,
|
|
"step": 1,
|
|
"tooltip": "Frames per second for video playback",
|
|
},
|
|
),
|
|
"max_width": (
|
|
"INT",
|
|
{
|
|
"default": 2048,
|
|
"tooltip": "The max width to use. Frames larger than the"
|
|
+ " specified value will be resized. Larger sizes may run"
|
|
+ " slower. Set to -1 for no resizing.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "view_video_pano"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "pytorch360convert"
|
|
|
|
def __init__(self) -> None:
|
|
self.type = "video_pano"
|
|
|
|
def view_video_pano(
|
|
self, video_frames: torch.Tensor, fps: int = 30, max_width: int = 2048
|
|
) -> Dict[str, Dict[str, str]]:
|
|
"""Process and display the panoramic video in the viewer.
|
|
|
|
This method handles the conversion of the input tensor frames to a viewable video format by:
|
|
1. Processing each frame in the batch
|
|
2. Converting each tensor to a numpy array
|
|
3. Converting float values to uint8 if necessary
|
|
4. Converting grayscale to RGB if necessary
|
|
5. Resizing the frames if they exceed max_width
|
|
6. Creating a base64-encoded video from the frames
|
|
|
|
Args:
|
|
video_frames (torch.Tensor): Input tensor containing the panoramic video frames.
|
|
Should be in format (B, H, W, C).
|
|
fps (int, optional): Frames per second for video playback. Default: 30
|
|
max_width (int, optional): Maximum width for resizing. Frames will not be
|
|
resized if they are smaller than the specified size. Set to -1 to
|
|
disable resizing. Default: 2048
|
|
|
|
Returns:
|
|
Dict[str, Dict[str, str]]: Dictionary containing the UI update information
|
|
with the base64-encoded video data and playback parameters.
|
|
"""
|
|
# Ensure we have batch dimension
|
|
if len(video_frames.shape) != 4:
|
|
raise ValueError("Expected video frames in batch format (B, H, W, C)")
|
|
|
|
# Create a list to store processed frames
|
|
processed_frames = []
|
|
|
|
# Process each frame
|
|
for frame in video_frames:
|
|
# Convert to numpy and proper format
|
|
frame_np = frame.cpu().numpy()
|
|
|
|
# Convert to uint8 if needed
|
|
if frame_np.dtype != np.uint8:
|
|
frame_np = (frame_np * 255).astype(np.uint8)
|
|
|
|
# Handle grayscale images
|
|
if len(frame_np.shape) == 2 or frame_np.shape[2] == 1:
|
|
frame_np = np.repeat(frame_np[..., np.newaxis], 3, axis=2)
|
|
|
|
# Convert to PIL Image
|
|
pil_frame = Image.fromarray(frame_np)
|
|
|
|
# Optionally resize image
|
|
if max_width > 0 and (
|
|
pil_frame.size[0] > max_width or pil_frame.size[1] > max_width
|
|
):
|
|
new_size = tuple(
|
|
[int(max_width * x / max(pil_frame.size)) for x in pil_frame.size]
|
|
)
|
|
pil_frame = pil_frame.resize(
|
|
new_size, resample=Image.Resampling.LANCZOS
|
|
)
|
|
|
|
# Add to processed frames
|
|
processed_frames.append(pil_frame)
|
|
|
|
# Check if we have any frames
|
|
if not processed_frames:
|
|
return {"ui": {"error": "No frames found in input"}}
|
|
|
|
# Create a list to store base64 strings of each frame
|
|
frame_data = []
|
|
|
|
# Convert each frame to base64
|
|
for frame in processed_frames:
|
|
buffered = BytesIO()
|
|
frame.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|
frame_data.append(f"data:image/png;base64,{img_str}")
|
|
|
|
# Return the frame data, count, and fps
|
|
return {
|
|
"ui": {
|
|
"pano_video_preview": frame_data[0], # First frame as preview
|
|
"pano_video_frames": frame_data, # All frames as list of strings
|
|
"frame_count": str(len(processed_frames)),
|
|
"fps": str(fps),
|
|
"video_type": "360_equirectangular",
|
|
}
|
|
}
|
|
|
|
|
|
# Update mappings to include the new node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PanoramaViewerNode": PanoramaViewerNode,
|
|
"PanoramaVideoViewerNode": PanoramaVideoViewerNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PanoramaViewerNode": "Preview 360 Panorama",
|
|
"PanoramaVideoViewerNode": "Preview 360 Video Panorama",
|
|
}
|