95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import torch
|
||
import numpy as np
|
||
import cv2
|
||
import math
|
||
|
||
from .projection import equirect_to_perspective, strip_to_equirect
|
||
|
||
|
||
class EquirectToStrip90:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"face_size": ("INT", {"default": 512, "min": 128, "max": 2048}),
|
||
"order": ("STRING", {"default": "3,2,1,0,4,5", "multiline": False}),
|
||
"flip_sides": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "STRING", "BOOLEAN")
|
||
RETURN_NAMES = ("strip", "order", "flip_sides")
|
||
FUNCTION = "convert"
|
||
CATEGORY = "360"
|
||
|
||
def convert(self, image, face_size, order, flip_sides):
|
||
img = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
||
|
||
# Parse order string to get indices
|
||
order_list = [int(x.strip()) for x in order.split(",")]
|
||
|
||
# All 6 views in standard order: front, left, back, right, top, bottom
|
||
all_views = [
|
||
("0_front", 0, 0),
|
||
("1_left", -math.pi/2, 0),
|
||
("2_back", math.pi, 0),
|
||
("3_right", math.pi/2, 0),
|
||
("4_top", 0, -math.pi/2),
|
||
("5_bottom", 0, math.pi/2),
|
||
]
|
||
|
||
faces = []
|
||
for i, idx in enumerate(order_list):
|
||
_, yaw, pitch = all_views[idx]
|
||
face = equirect_to_perspective(img, yaw, pitch, 90, face_size)
|
||
# Optionally flip first 4 faces (horizontal views) vertically
|
||
if flip_sides and idx < 4:
|
||
face = cv2.flip(face, 0)
|
||
faces.append(face)
|
||
|
||
strip = np.concatenate(faces, axis=1)
|
||
strip = torch.from_numpy(strip.astype(np.float32) / 255.0).unsqueeze(0)
|
||
|
||
return (strip, order, flip_sides)
|
||
|
||
|
||
class Strip90ToEquirect:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"strip": ("IMAGE",),
|
||
"width": ("INT", {"default": 2048}),
|
||
"height": ("INT", {"default": 1024}),
|
||
"face_size": ("INT", {"default": 512}),
|
||
"order": ("STRING", {"default": "0,1,2,3,4,5", "multiline": False}),
|
||
"flip_sides": ("BOOLEAN", {"default": False}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "convert"
|
||
CATEGORY = "360"
|
||
|
||
def convert(self, strip, width, height, face_size, order, flip_sides):
|
||
strip_img = (strip[0].cpu().numpy() * 255).astype(np.uint8)
|
||
order_list = [int(x.strip()) for x in order.split(",")]
|
||
out = strip_to_equirect(strip_img, width, height, face_size, order_list, flip_sides)
|
||
out = torch.from_numpy(out.astype(np.float32) / 255.0).unsqueeze(0)
|
||
return (out,)
|
||
|
||
# -------------------------
|
||
# ComfyUI registration
|
||
# -------------------------
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"EquirectToStrip90": EquirectToStrip90,
|
||
"Strip90ToEquirect": Strip90ToEquirect,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"EquirectToStrip90": "360 → 6×90° Strip",
|
||
"Strip90ToEquirect": "6×90° Strip → 360",
|
||
}
|