383 lines
14 KiB
Python
383 lines
14 KiB
Python
"""
|
||
ComfyUI-CameraPack: Convert camera parameters between formats.
|
||
|
||
Converts DepthAnythingV3 extrinsics/intrinsics to ComfyUI's Load3DCamera format.
|
||
"""
|
||
|
||
import ast
|
||
import re
|
||
import torch
|
||
import numpy as np
|
||
|
||
|
||
class DA3ToLoad3DCamera:
|
||
"""
|
||
Convert DepthAnythingV3 camera parameters to ComfyUI Load3DCamera format.
|
||
|
||
DA3 outputs:
|
||
- extrinsics: [B, 4, 4] pose matrix (rotation + translation)
|
||
- intrinsics: [B, 3, 3] matrix [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
|
||
|
||
Load3DCamera expects:
|
||
- position: {x, y, z} camera position in world space
|
||
- target: {x, y, z} look-at point
|
||
- zoom: int (derived from focal length)
|
||
- cameraType: "perspective" or "orthographic"
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"extrinsics": ("EXTRINSICS",),
|
||
"intrinsics": ("INTRINSICS",),
|
||
},
|
||
"optional": {
|
||
"target_distance": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.1,
|
||
"max": 100.0,
|
||
"step": 0.1,
|
||
"tooltip": "Distance from camera to target point along view direction"
|
||
}),
|
||
"zoom_scale": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.1,
|
||
"max": 10.0,
|
||
"step": 0.1,
|
||
"tooltip": "Scale factor for zoom (multiply focal length)"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LOAD3D_CAMERA",)
|
||
RETURN_NAMES = ("camera_info",)
|
||
FUNCTION = "convert"
|
||
CATEGORY = "3d/camera"
|
||
|
||
def convert(self, extrinsics, intrinsics, target_distance=1.0, zoom_scale=1.0):
|
||
"""
|
||
Convert DA3 camera parameters to Load3DCamera format.
|
||
|
||
The extrinsics matrix is a 4x4 world-to-camera transform:
|
||
[R | t]
|
||
[0 | 1]
|
||
|
||
Camera position in world space: pos = -R^T @ t
|
||
Camera forward direction: forward = R^T @ [0, 0, 1]
|
||
"""
|
||
# Handle batch dimension - take first item
|
||
if isinstance(extrinsics, torch.Tensor):
|
||
ext = extrinsics[0].cpu().numpy() if extrinsics.dim() > 2 else extrinsics.cpu().numpy()
|
||
else:
|
||
ext = np.array(extrinsics)
|
||
|
||
if isinstance(intrinsics, torch.Tensor):
|
||
intr = intrinsics[0].cpu().numpy() if intrinsics.dim() > 2 else intrinsics.cpu().numpy()
|
||
else:
|
||
intr = np.array(intrinsics)
|
||
|
||
# Flatten and handle different extrinsics formats
|
||
ext = ext.flatten()
|
||
if ext.size == 16:
|
||
ext = ext.reshape(4, 4)
|
||
elif ext.size == 12:
|
||
# 3x4 matrix - add [0, 0, 0, 1] row to make 4x4
|
||
ext = ext.reshape(3, 4)
|
||
ext = np.vstack([ext, [0, 0, 0, 1]])
|
||
else:
|
||
raise ValueError(f"Unexpected extrinsics size: {ext.size}, expected 12 or 16")
|
||
|
||
# Ensure intrinsics is 3x3
|
||
intr = intr.flatten()
|
||
if intr.size == 9:
|
||
intr = intr.reshape(3, 3)
|
||
else:
|
||
raise ValueError(f"Unexpected intrinsics size: {intr.size}, expected 9")
|
||
|
||
# Extract rotation and translation from extrinsics
|
||
R = ext[:3, :3] # 3x3 rotation matrix
|
||
t = ext[:3, 3] # 3x1 translation vector
|
||
|
||
# Camera position in world space: pos = -R^T @ t
|
||
R_T = R.T
|
||
camera_pos = -R_T @ t
|
||
|
||
# Camera forward direction (Z-axis in camera space transformed to world)
|
||
# In OpenGL/standard convention, camera looks down -Z, so forward is [0, 0, -1]
|
||
# In some conventions it's [0, 0, 1]. We'll use [0, 0, 1] and let user adjust.
|
||
forward = R_T @ np.array([0, 0, 1])
|
||
|
||
# Target point is camera position + forward * distance
|
||
target_pos = camera_pos + forward * target_distance
|
||
|
||
# Extract focal length from intrinsics
|
||
fx = float(intr[0, 0])
|
||
fy = float(intr[1, 1])
|
||
focal_length = (fx + fy) / 2.0
|
||
|
||
# Compute zoom - this is somewhat arbitrary, scale to reasonable range
|
||
zoom = int(focal_length * zoom_scale)
|
||
zoom = max(1, min(zoom, 1000)) # Clamp to reasonable range
|
||
|
||
# Build camera_info dict
|
||
camera_info = {
|
||
"position": {
|
||
"x": float(camera_pos[0]),
|
||
"y": float(camera_pos[1]),
|
||
"z": float(camera_pos[2]),
|
||
},
|
||
"target": {
|
||
"x": float(target_pos[0]),
|
||
"y": float(target_pos[1]),
|
||
"z": float(target_pos[2]),
|
||
},
|
||
"zoom": zoom,
|
||
"cameraType": "perspective",
|
||
}
|
||
|
||
return (camera_info,)
|
||
|
||
|
||
class CreateLoad3DCamera:
|
||
"""
|
||
Manually create a Load3DCamera from position, target, and zoom values.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"pos_x": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"pos_y": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"pos_z": ("FLOAT", {"default": 5.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"target_x": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"target_y": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"target_z": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.1}),
|
||
"zoom": ("INT", {"default": 50, "min": 1, "max": 1000, "step": 1}),
|
||
"camera_type": (["perspective", "orthographic"], {"default": "perspective"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LOAD3D_CAMERA",)
|
||
RETURN_NAMES = ("camera_info",)
|
||
FUNCTION = "create"
|
||
CATEGORY = "3d/camera"
|
||
|
||
def create(self, pos_x, pos_y, pos_z, target_x, target_y, target_z, zoom, camera_type):
|
||
camera_info = {
|
||
"position": {"x": pos_x, "y": pos_y, "z": pos_z},
|
||
"target": {"x": target_x, "y": target_y, "z": target_z},
|
||
"zoom": zoom,
|
||
"cameraType": camera_type,
|
||
}
|
||
return (camera_info,)
|
||
|
||
|
||
class CameraIntrinsics:
|
||
"""
|
||
Create an INTRINSICS tensor from a pasted 3x3 matrix.
|
||
|
||
Accepts formats like:
|
||
[[3.2782e+03, 0, 960], [0, 2.4555e+03, 540], [0, 0, 1]]
|
||
tensor([[[3278.2, 0, 960], [0, 2455.5, 540], [0, 0, 1]]])
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"matrix": ("STRING", {
|
||
"default": "[[1000, 0, 960],\n [0, 1000, 540],\n [0, 0, 1]]",
|
||
"multiline": True,
|
||
"tooltip": "3x3 intrinsics matrix [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("INTRINSICS",)
|
||
RETURN_NAMES = ("intrinsics",)
|
||
FUNCTION = "create"
|
||
CATEGORY = "3d/camera"
|
||
|
||
def create(self, matrix):
|
||
text = matrix.strip()
|
||
# Strip tensor(...) wrapper if pasted from PyTorch output
|
||
text = re.sub(r'^tensor\(', '', text)
|
||
text = re.sub(r'\)\s*$', '', text)
|
||
text = text.strip()
|
||
|
||
parsed = ast.literal_eval(text)
|
||
arr = np.array(parsed, dtype=np.float32)
|
||
|
||
# Handle [1, 3, 3] input
|
||
if arr.ndim == 3 and arr.shape == (1, 3, 3):
|
||
arr = arr[0]
|
||
|
||
if arr.shape != (3, 3):
|
||
raise ValueError(f"Expected 3x3 matrix, got shape {arr.shape}")
|
||
|
||
intrinsics = torch.from_numpy(arr).unsqueeze(0) # [1, 3, 3]
|
||
return (intrinsics,)
|
||
|
||
|
||
class GetExtrinsicsFromBatch:
|
||
"""Slice a sub-batch (or single) of extrinsics from a batched [N, 4, 4]
|
||
EXTRINSICS tensor — mirrors ComfyUI's native `ImageFromBatch`.
|
||
|
||
Use to pick one camera out of a 12-face panorama split, the first view
|
||
of an HYWM2 reconstruction, etc. Negative `batch_index` wraps from the
|
||
end; `length=0` returns the full tail starting at `batch_index`.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"extrinsics": ("EXTRINSICS", {
|
||
"tooltip": "Batched w2c [N, 4, 4] (or [1, N, 4, 4]).",
|
||
}),
|
||
"batch_index": ("INT", {
|
||
"default": 0, "min": -4096, "max": 4096, "step": 1,
|
||
"tooltip": "Starting index. Negative wraps from the end "
|
||
"(e.g. -1 = last extrinsic).",
|
||
}),
|
||
"length": ("INT", {
|
||
"default": 1, "min": 0, "max": 4096, "step": 1,
|
||
"tooltip": "Number of consecutive extrinsics to return "
|
||
"starting at batch_index. 0 = take everything "
|
||
"from batch_index to the end. Default 1 = "
|
||
"single extrinsic.",
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("EXTRINSICS",)
|
||
RETURN_NAMES = ("extrinsics",)
|
||
FUNCTION = "slice_batch"
|
||
CATEGORY = "3d/camera"
|
||
|
||
def slice_batch(self, extrinsics, batch_index: int = 0, length: int = 1):
|
||
if not isinstance(extrinsics, torch.Tensor):
|
||
extrinsics = torch.as_tensor(extrinsics).float()
|
||
ext = extrinsics
|
||
if ext.dim() == 4 and ext.shape[0] == 1:
|
||
ext = ext[0]
|
||
if ext.dim() != 3 or ext.shape[-2:] not in ((4, 4), (3, 4)):
|
||
raise ValueError(
|
||
f"GetExtrinsicsFromBatch: expected [N, 4, 4] or [N, 3, 4]; "
|
||
f"got {tuple(extrinsics.shape)}"
|
||
)
|
||
N = int(ext.shape[0])
|
||
# Negative wrap.
|
||
start = batch_index % N if N > 0 else 0
|
||
if length <= 0:
|
||
end = N
|
||
else:
|
||
end = min(start + length, N)
|
||
if start >= end:
|
||
raise ValueError(
|
||
f"GetExtrinsicsFromBatch: empty slice (start={start} >= "
|
||
f"end={end}, N={N}, batch_index={batch_index}, length={length})"
|
||
)
|
||
# `.clone()` (not just `.contiguous()`) — comfy-env's shm IPC
|
||
# serializer reads the underlying storage size, not the view's
|
||
# logical size. A sliced view of a 12×4×4 tensor still points at
|
||
# the 192-element backing buffer; the receiver tries to reshape
|
||
# 192 floats into [1, 4, 4] and explodes. `.clone()` allocates
|
||
# fresh, view-free storage matching the slice's shape exactly.
|
||
sliced = ext[start:end].clone().contiguous()
|
||
# Squeeze single-element batches so single-camera consumers
|
||
# (SHARP's predict, etc.) get the [4, 4] they expect rather than
|
||
# [1, 4, 4]. Multi-element slices stay batched.
|
||
if sliced.shape[0] == 1:
|
||
sliced = sliced[0].contiguous()
|
||
return (sliced,)
|
||
|
||
|
||
class GetIntrinsicsFromBatch:
|
||
"""Slice a sub-batch (or single) of intrinsics from a batched [N, 3, 3]
|
||
INTRINSICS tensor — mirrors ComfyUI's native `ImageFromBatch`.
|
||
|
||
Same semantics as `GetExtrinsicsFromBatch` but for K matrices.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"intrinsics": ("INTRINSICS", {
|
||
"tooltip": "Batched K [N, 3, 3] (or [1, N, 3, 3]).",
|
||
}),
|
||
"batch_index": ("INT", {
|
||
"default": 0, "min": -4096, "max": 4096, "step": 1,
|
||
"tooltip": "Starting index. Negative wraps from the end.",
|
||
}),
|
||
"length": ("INT", {
|
||
"default": 1, "min": 0, "max": 4096, "step": 1,
|
||
"tooltip": "Number of consecutive intrinsics to return. "
|
||
"0 = take everything from batch_index to the "
|
||
"end. Default 1 = single K matrix.",
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("INTRINSICS",)
|
||
RETURN_NAMES = ("intrinsics",)
|
||
FUNCTION = "slice_batch"
|
||
CATEGORY = "3d/camera"
|
||
|
||
def slice_batch(self, intrinsics, batch_index: int = 0, length: int = 1):
|
||
if not isinstance(intrinsics, torch.Tensor):
|
||
intrinsics = torch.as_tensor(intrinsics).float()
|
||
intr = intrinsics
|
||
if intr.dim() == 4 and intr.shape[0] == 1:
|
||
intr = intr[0]
|
||
if intr.dim() == 2:
|
||
# Single K passed unbatched — wrap then slice.
|
||
intr = intr.unsqueeze(0)
|
||
if intr.dim() != 3 or intr.shape[-2:] != (3, 3):
|
||
raise ValueError(
|
||
f"GetIntrinsicsFromBatch: expected [N, 3, 3]; "
|
||
f"got {tuple(intrinsics.shape)}"
|
||
)
|
||
N = int(intr.shape[0])
|
||
start = batch_index % N if N > 0 else 0
|
||
if length <= 0:
|
||
end = N
|
||
else:
|
||
end = min(start + length, N)
|
||
if start >= end:
|
||
raise ValueError(
|
||
f"GetIntrinsicsFromBatch: empty slice (start={start} >= "
|
||
f"end={end}, N={N}, batch_index={batch_index}, length={length})"
|
||
)
|
||
# See GetExtrinsicsFromBatch.slice_batch for why `.clone()` (not
|
||
# just `.contiguous()`) — comfy-env's shm IPC keeps the original
|
||
# storage when the tensor is a view.
|
||
sliced = intr[start:end].clone().contiguous()
|
||
# Squeeze single-element batches so single-camera consumers
|
||
# (SHARP's predict — `intrinsics[0, 0].item()` for fx, etc.)
|
||
# get the [3, 3] they expect rather than [1, 3, 3].
|
||
if sliced.shape[0] == 1:
|
||
sliced = sliced[0].contiguous()
|
||
return (sliced,)
|
||
|
||
|
||
# Node registration
|
||
NODE_CLASS_MAPPINGS = {
|
||
"DA3ToLoad3DCamera": DA3ToLoad3DCamera,
|
||
"CreateLoad3DCamera": CreateLoad3DCamera,
|
||
"CameraIntrinsics": CameraIntrinsics,
|
||
"GetExtrinsicsFromBatch": GetExtrinsicsFromBatch,
|
||
"GetIntrinsicsFromBatch": GetIntrinsicsFromBatch,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"DA3ToLoad3DCamera": "DA3 to Load3D Camera",
|
||
"CreateLoad3DCamera": "Create Load3D Camera",
|
||
"CameraIntrinsics": "Camera Intrinsics",
|
||
"GetExtrinsicsFromBatch": "Get Extrinsics From Batch",
|
||
"GetIntrinsicsFromBatch": "Get Intrinsics From Batch",
|
||
}
|