130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
import torch
|
|
import numpy as np
|
|
import cv2
|
|
import math
|
|
|
|
|
|
class DAP_ERP_to_Cubemap:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"face_size": (
|
|
"INT",
|
|
{"default": 1024, "min": 256, "max": 4096, "step": 64},
|
|
),
|
|
"layout": (["six_faces", "cross"], {"default": "six_faces"}),
|
|
"interpolation": (
|
|
["lanczos", "cubic", "linear", "nearest"],
|
|
{"default": "lanczos"},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "DAP/Conversions"
|
|
|
|
def convert(self, image, face_size, layout, interpolation):
|
|
# ComfyUI image is [B, H, W, C], 0..1
|
|
B, H, W, C = image.shape
|
|
|
|
interp_map = {
|
|
"nearest": cv2.INTER_NEAREST,
|
|
"linear": cv2.INTER_LINEAR,
|
|
"cubic": cv2.INTER_CUBIC,
|
|
"lanczos": cv2.INTER_LANCZOS4,
|
|
}
|
|
interp = interp_map[interpolation]
|
|
|
|
faces_order = ["right", "left", "top", "bottom", "front", "back"]
|
|
|
|
all_results = []
|
|
|
|
for b in range(B):
|
|
erp_img = (image[b].cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
faces = {}
|
|
for face_name in faces_order:
|
|
map_x_norm, map_y_norm = self.build_face_map(face_size, face_name)
|
|
|
|
# Apply map
|
|
map_x = map_x_norm * (W - 1)
|
|
map_y = map_y_norm * (H - 1)
|
|
|
|
face_img = cv2.remap(
|
|
erp_img,
|
|
map_x,
|
|
map_y,
|
|
interpolation=interp,
|
|
borderMode=cv2.BORDER_WRAP,
|
|
)
|
|
faces[face_name] = face_img
|
|
|
|
if layout == "six_faces":
|
|
# Return as a batch of 6 images
|
|
for name in faces_order:
|
|
res = torch.from_numpy(faces[name].astype(np.float32) / 255.0)
|
|
all_results.append(res)
|
|
else:
|
|
# Return as a single cross layout image (3x4 grid)
|
|
cross = self.make_cross_layout(faces, face_size, C)
|
|
res = torch.from_numpy(cross.astype(np.float32) / 255.0)
|
|
all_results.append(res)
|
|
|
|
final_image = torch.stack(all_results)
|
|
return (final_image,)
|
|
|
|
def build_face_map(self, face_size, face):
|
|
s = face_size
|
|
jj, ii = np.meshgrid(
|
|
np.arange(s, dtype=np.float32), np.arange(s, dtype=np.float32)
|
|
)
|
|
a = 2.0 * (jj + 0.5) / s - 1.0
|
|
b = 2.0 * (ii + 0.5) / s - 1.0
|
|
|
|
if face == "right": # +X
|
|
dx, dy, dz = np.ones_like(a), -b, -a
|
|
elif face == "left": # -X
|
|
dx, dy, dz = -np.ones_like(a), -b, a
|
|
elif face == "top": # +Y
|
|
dx, dy, dz = a, np.ones_like(a), b
|
|
elif face == "bottom": # -Y
|
|
dx, dy, dz = a, -np.ones_like(a), -b
|
|
elif face == "front": # +Z
|
|
dx, dy, dz = a, -b, np.ones_like(a)
|
|
elif face == "back": # -Z
|
|
dx, dy, dz = -a, -b, -np.ones_like(a)
|
|
else:
|
|
raise ValueError(f"Unknown face: {face}")
|
|
|
|
norm = np.sqrt(dx * dx + dy * dy + dz * dz)
|
|
dx /= norm
|
|
dy /= norm
|
|
dz /= norm
|
|
|
|
theta = np.arctan2(dz, dx)
|
|
phi = np.arcsin(dy)
|
|
|
|
map_x_norm = (theta + math.pi) / (2.0 * math.pi)
|
|
map_y_norm = (math.pi / 2 - phi) / math.pi
|
|
|
|
return map_x_norm.astype(np.float32), map_y_norm.astype(np.float32)
|
|
|
|
def make_cross_layout(self, faces, face_size, channels):
|
|
S = face_size
|
|
canvas = np.zeros((3 * S, 4 * S, channels), dtype=np.uint8)
|
|
|
|
def put(name, row, col):
|
|
canvas[row * S : (row + 1) * S, col * S : (col + 1) * S] = faces[name]
|
|
|
|
put("top", 0, 1)
|
|
put("left", 1, 0)
|
|
put("front", 1, 1)
|
|
put("right", 1, 2)
|
|
put("back", 1, 3)
|
|
put("bottom", 2, 1)
|
|
return canvas
|