Add video panorama node (#7)

This commit is contained in:
jpgallegoar
2025-05-25 13:25:57 -06:00
committed by GitHub
parent 7c082052fc
commit 5b1dd037ba
2 changed files with 466 additions and 9 deletions
+307 -7
View File
@@ -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
};
}
}
});
+159 -2
View File
@@ -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",
}