From 5b1dd037ba693e837233246bbfaebdab8933f45f Mon Sep 17 00:00:00 2001 From: jpgallegoar <65631402+jpgallegoar@users.noreply.github.com> Date: Sun, 25 May 2025 21:25:57 +0200 Subject: [PATCH] Add video panorama node (#7) --- js/pano_viewer.js | 314 ++++++++++++++++++++++++++++++++++++++++++++-- nodes.py | 161 +++++++++++++++++++++++- 2 files changed, 466 insertions(+), 9 deletions(-) diff --git a/js/pano_viewer.js b/js/pano_viewer.js index 7e90186..e02dbcc 100644 --- a/js/pano_viewer.js +++ b/js/pano_viewer.js @@ -43,14 +43,23 @@ class MinimalPanoViewer { loadImage(dataUrl) { const loader = new THREE.TextureLoader(); - loader.load(dataUrl, (texture) => { - texture.colorSpace = THREE.SRGBColorSpace; - texture.mapping = THREE.EquirectangularReflectionMapping; - texture.minFilter = texture.magFilter = THREE.LinearFilter; + // Add error handling + loader.load( + dataUrl, + (texture) => { + console.log("Texture loaded successfully"); + texture.colorSpace = THREE.SRGBColorSpace; + texture.mapping = THREE.EquirectangularReflectionMapping; + texture.minFilter = texture.magFilter = THREE.LinearFilter; - // Set the texture as the background of the scene - this.scene.background = texture; - }); + // Set the texture as the background of the scene + this.scene.background = texture; + }, + undefined, + (error) => { + console.error("Error loading texture:", error); + } + ); } onMouseDown(event) { @@ -100,6 +109,212 @@ class MinimalPanoViewer { } } +// Extended viewer for video panoramas +class MinimalPanoVideoViewer extends MinimalPanoViewer { + constructor(container) { + super(container); + + // Video-specific properties + this.frameCount = 0; + this.fps = 30; + this.currentFrame = 0; + this.isPlaying = false; + this.lastFrameTime = 0; + this.frames = []; // Store all frame data URLs + this.isLoadingFrame = false; // Track if a frame is currently loading + + // Comment out this line to remove controls completely + // this.createVideoControls(container); + + // Create a dummy controls container and elements to avoid errors + this.controlsContainer = document.createElement("div"); + this.controlsContainer.style.display = "none"; + this.playButton = document.createElement("button"); + this.timelineSlider = document.createElement("input"); + this.frameCounter = document.createElement("span"); + + // Add status indicator + this.statusText = document.createElement("div"); + this.statusText.style.position = "absolute"; + this.statusText.style.top = "10px"; + this.statusText.style.left = "10px"; + this.statusText.style.color = "white"; + this.statusText.style.backgroundColor = "rgba(0, 0, 0, 0.5)"; + this.statusText.style.padding = "5px"; + this.statusText.style.borderRadius = "5px"; + this.statusText.style.fontSize = "12px"; + this.statusText.innerText = "Loading..."; + this.statusText.style.display = "none"; + container.appendChild(this.statusText); + } + + createVideoControls(container) { + // Create controls container + this.controlsContainer = document.createElement("div"); + this.controlsContainer.style.position = "absolute"; + this.controlsContainer.style.bottom = "10px"; + this.controlsContainer.style.left = "10px"; + this.controlsContainer.style.right = "10px"; + this.controlsContainer.style.display = "flex"; + this.controlsContainer.style.alignItems = "center"; + this.controlsContainer.style.justifyContent = "center"; + this.controlsContainer.style.backgroundColor = "rgba(0, 0, 0, 0.5)"; + this.controlsContainer.style.padding = "5px"; + this.controlsContainer.style.borderRadius = "5px"; + container.appendChild(this.controlsContainer); + + // Play/Pause button + this.playButton = document.createElement("button"); + this.playButton.innerText = "▶"; + this.playButton.style.marginRight = "5px"; + this.playButton.style.width = "30px"; + this.playButton.style.height = "30px"; + this.playButton.style.borderRadius = "15px"; + this.playButton.style.border = "none"; + this.playButton.style.background = "#444"; + this.playButton.style.color = "white"; + this.playButton.style.cursor = "pointer"; + this.playButton.addEventListener("click", this.togglePlay.bind(this)); + this.controlsContainer.appendChild(this.playButton); + + // Timeline slider + this.timelineSlider = document.createElement("input"); + this.timelineSlider.type = "range"; + this.timelineSlider.min = "0"; + this.timelineSlider.max = "100"; + this.timelineSlider.value = "0"; + this.timelineSlider.style.flex = "1"; + this.timelineSlider.style.margin = "0 10px"; + this.timelineSlider.addEventListener("input", this.onTimelineChange.bind(this)); + this.controlsContainer.appendChild(this.timelineSlider); + + // Frame counter + this.frameCounter = document.createElement("span"); + this.frameCounter.style.color = "white"; + this.frameCounter.style.fontSize = "12px"; + this.frameCounter.style.marginLeft = "5px"; + this.frameCounter.innerText = "0/0"; + this.controlsContainer.appendChild(this.frameCounter); + + // Hide controls initially + this.controlsContainer.style.display = "none"; + } + + togglePlay() { + this.isPlaying = !this.isPlaying; + this.playButton.innerText = this.isPlaying ? "⏸" : "▶"; + + if (this.isPlaying) { + this.lastFrameTime = performance.now(); + this.playFrames(); + } + } + + onTimelineChange(event) { + const frameIndex = Math.floor((parseInt(event.target.value) / 100) * (this.frameCount - 1)); + this.currentFrame = frameIndex; + this.updateFrameCounter(); + // Load the frame at the current position + this.loadCurrentFrame(); + } + + loadCurrentFrame() { + if (this.isLoadingFrame) { + console.log("Already loading a frame, skipping"); + return; + } + + if (this.frames.length > 0 && this.currentFrame < this.frames.length) { + console.log(`Loading frame ${this.currentFrame + 1}/${this.frameCount}`); + this.statusText.innerText = `Loading frame ${this.currentFrame + 1}/${this.frameCount}`; + this.statusText.style.display = "block"; + + this.isLoadingFrame = true; + + // Use setTimeout to ensure the UI updates before loading the texture + setTimeout(() => { + this.loadImage(this.frames[this.currentFrame]); + this.isLoadingFrame = false; + + // Hide status after a short delay + setTimeout(() => { + this.statusText.style.display = "none"; + }, 500); + }, 10); + } else { + console.warn(`Cannot load frame ${this.currentFrame + 1}, total frames: ${this.frames.length}`); + this.statusText.innerText = "Error loading frame"; + this.statusText.style.display = "block"; + } + } + + updateFrameCounter() { + this.frameCounter.innerText = `${this.currentFrame + 1}/${this.frameCount}`; + this.timelineSlider.value = ((this.currentFrame / (this.frameCount - 1)) * 100).toString(); + } + + playFrames() { + if (!this.isPlaying) return; + + const now = performance.now(); + const elapsed = now - this.lastFrameTime; + const frameTime = 1000 / this.fps; + + if (elapsed >= frameTime && !this.isLoadingFrame) { + this.lastFrameTime = now - (elapsed % frameTime); + this.currentFrame = (this.currentFrame + 1) % this.frameCount; + this.updateFrameCounter(); + // Load the next frame + this.loadCurrentFrame(); + } + + requestAnimationFrame(this.playFrames.bind(this)); + } + + setVideoData(frames, frameCount, fps) { + console.log("Setting video data with", frames.length, "frames"); + this.frames = frames; + this.frameCount = parseInt(frameCount); + this.fps = parseInt(fps); + this.currentFrame = 0; + + // Update UI + this.timelineSlider.max = Math.max(1, this.frameCount - 1).toString(); + this.updateFrameCounter(); + + // Show controls if we have valid video data + if (this.frameCount > 1) { + // Auto-start playback immediately + this.isPlaying = true; + this.playButton.innerText = "⏸"; + this.lastFrameTime = performance.now(); + this.playFrames(); + + // Either hide controls or show them based on preference + this.controlsContainer.style.display = "none"; // Hide controls + // this.controlsContainer.style.display = "flex"; // Show controls + + // Preload the first frame immediately + this.loadCurrentFrame(); + } else { + console.warn("Invalid frame count:", this.frameCount); + } + + // Status update + this.statusText.innerText = `Loaded ${this.frameCount} frames at ${this.fps} FPS`; + this.statusText.style.display = "block"; + setTimeout(() => { + this.statusText.style.display = "none"; + }, 3000); + } + + // Override the update method to handle video frames + update() { + super.update(); + // Additional video-specific updates would go here + } +} + app.registerExtension({ name: "Pano.Viewer", async beforeRegisterNodeDef(nodeType, nodeData, app) { @@ -113,6 +328,7 @@ app.registerExtension({ container.style.backgroundColor = "#000000"; container.style.borderRadius = "8px"; container.style.overflow = "hidden"; + container.style.position = "relative"; // Add position relative for controls positioning // Initialize the widget let panoramaWidget = this.addDOMWidget("panoramapreview", "preview", container, { @@ -143,5 +359,89 @@ app.registerExtension({ return ret }; } + else if (nodeData.name === "PanoramaVideoViewerNode") { + // Override the node's onCreate + let originalOnCreated = nodeType.prototype.onNodeCreated + nodeType.prototype.onNodeCreated = function() { + let ret = originalOnCreated?.apply?.(this, arguments) + // Create viewer container + const container = document.createElement("div"); + container.style.backgroundColor = "#000000"; + container.style.borderRadius = "8px"; + container.style.overflow = "hidden"; + container.style.position = "relative"; // Add position relative for controls positioning + + // Initialize the widget + let panoramaWidget = this.addDOMWidget("panoramavideopreview", "preview", container, { + serialize: false, hideOnZoom: false + }); + let node = this + panoramaWidget.computeSize = function(width) { + let height = node.size[0] - 10; + this.computedHeight = height + 10; + return [width, height]; + } + + // Initialize video viewer + this.viewer = new MinimalPanoVideoViewer(container); + + panoramaWidget.options.afterResize = this.viewer.resizeView + requestAnimationFrame(this.viewer.resizeView) + return ret + }; + + // Override the node's onExecute + let originalOnExecuted = nodeType.prototype.onExecuted + nodeType.prototype.onExecuted = function(output) { + let ret = originalOnExecuted?.apply?.(this, arguments) + console.log("PanoramaVideoViewerNode output:", output); + + if (output?.pano_video_preview) { + // Load the first frame preview + console.log("Loading preview frame"); + this.viewer?.loadImage(output.pano_video_preview.join('')); + + // Set video data if available + if (output?.pano_video_frames && output?.frame_count && output?.fps) { + console.log("Setting video data"); + console.log("Frame count:", output.frame_count.join('')); + console.log("FPS:", output.fps.join('')); + + // Check if pano_video_frames is an array + if (Array.isArray(output.pano_video_frames)) { + const frames = []; + console.log("Number of frames received:", output.pano_video_frames.length); + + // Process frame data URLs - handle both string and array cases + for (let frameData of output.pano_video_frames) { + if (Array.isArray(frameData)) { + frames.push(frameData.join('')); + } else { + frames.push(frameData); + } + } + + console.log("Processed frames:", frames.length, "first frame type:", typeof frames[0]); + + // Use setTimeout to ensure UI updates properly before loading frames + setTimeout(() => { + this.viewer?.setVideoData( + frames, + output.frame_count.join(''), + output.fps.join('') + ); + }, 100); + } else { + console.error("Expected array for pano_video_frames, got:", typeof output.pano_video_frames); + } + } else { + console.warn("No pano_video_frames found in output"); + } + } else { + console.warn("No pano_video_preview found in output"); + } + return ret + }; + } } }); diff --git a/nodes.py b/nodes.py index 31c0713..a174697 100644 --- a/nodes.py +++ b/nodes.py @@ -119,6 +119,163 @@ class PanoramaViewerNode: return {"ui": {"pano_image": f"data:image/png;base64,{img_str}"}} -NODE_CLASS_MAPPINGS = {"PanoramaViewerNode": PanoramaViewerNode} +class PanoramaVideoViewerNode: + """A ComfyUI node that provides an interactive 360-degree panorama viewer for + equirectangular video content. -NODE_DISPLAY_NAME_MAPPINGS = {"PanoramaViewerNode": "Preview 360 Panorama"} + 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", +}