422 lines
13 KiB
Python
422 lines
13 KiB
Python
import numpy as np
|
|
import equilib as py360
|
|
import torch
|
|
from .distort import plane_to_equirect
|
|
from .misc import *
|
|
import math
|
|
import cv2
|
|
|
|
def bitwise_and_masks(mask1, mask2):
|
|
mask1 = mask1.cpu()
|
|
mask2 = mask2.cpu()
|
|
cv2_mask1 = np.array(mask1)
|
|
cv2_mask2 = np.array(mask2)
|
|
|
|
if cv2_mask1.shape == cv2_mask2.shape:
|
|
cv2_mask = cv2.bitwise_and(cv2_mask1, cv2_mask2)
|
|
return torch.from_numpy(cv2_mask)
|
|
else:
|
|
# do nothing - incompatible mask shape: mostly empty mask
|
|
print("incompatible mask shape: mostly empty mask")
|
|
return mask1
|
|
def calculate_equ_heght(heght):
|
|
# 计算 sqrt(x * x * 3)
|
|
result = math.sqrt(heght * heght * 3)
|
|
|
|
# 将结果四舍五入到最接近的64的倍数
|
|
result = int((result / 64 + 0.5)) * 64
|
|
|
|
if result < 512:
|
|
result = 512
|
|
|
|
return result
|
|
|
|
def calculate_face_heght(heght):
|
|
# 计算 sqrt(x * x / 3)
|
|
result = math.sqrt(heght * heght / 3)
|
|
|
|
# 将结果四舍五入到最接近的64的倍数
|
|
result = int((result / 64 + 0.5)) * 64
|
|
|
|
if result < 512:
|
|
result = 512
|
|
|
|
return result
|
|
|
|
# Equirectangular全景图，Cubemap 立方体贴图 Perspective 透视图
|
|
class PanoImageCube2Equ:
|
|
def __init__(self):
|
|
pass
|
|
|
|
CATEGORY = "PanoCard/convert"
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"height": ("INT", {
|
|
"default": 832,
|
|
"min": 0,
|
|
"step": 8
|
|
}),
|
|
"mask_scale": ("FLOAT", {
|
|
"default": 0,
|
|
"max": 1.0,
|
|
"min": 0,
|
|
"step": 0.01
|
|
}),
|
|
"mask_split": ("BOOLEAN", {"default": True, "label_on": "true", "label_off": "false"}),
|
|
},
|
|
"optional": {
|
|
"face": ("IMAGE",),
|
|
"long": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","MASK","MASK")
|
|
RETURN_NAMES = ("equ","equ_mask","equ_masks")
|
|
OUTPUT_IS_LIST = (False, False, True)
|
|
FUNCTION = "node"
|
|
|
|
def node(self, height, mask_scale, mask_split, face = None, long = None):
|
|
sub_height = 0
|
|
scale = mask_scale
|
|
result = np.array([])
|
|
order = "nearest"
|
|
if long is not None:
|
|
sub_height, sub_width = long.shape[1], long.shape[2]
|
|
|
|
if sub_height*6 != sub_width:
|
|
raise ValueError("传入图像宽高必须为6:1")
|
|
if height == 0:
|
|
height = calculate_equ_heght(sub_height)
|
|
|
|
long_face = long.permute(0, 3, 1, 2)
|
|
equ = py360.cube2equi(
|
|
cubemap = long_face,
|
|
height = height,
|
|
width = height * 2,
|
|
mode = order,
|
|
cube_format = 'horizon')
|
|
|
|
result = equ.permute(1, 2, 0).unsqueeze(0)
|
|
|
|
elif face is not None:
|
|
if not all(img.shape == face[0].shape for img in face):
|
|
raise ValueError("所有图像的尺寸必须相同")
|
|
|
|
sub_height = face[0].shape[0]
|
|
if height == 0:
|
|
height = calculate_equ_heght(sub_height)
|
|
|
|
face = tuple(face.permute(0, 3, 1, 2))
|
|
long_face = torch.cat(face, dim=2)
|
|
equ = py360.cube2equi(
|
|
cubemap = long_face,
|
|
height = height,
|
|
width = height * 2,
|
|
mode = order,
|
|
cube_format = 'horizon')
|
|
|
|
result = equ.permute(1, 2, 0).unsqueeze(0)
|
|
|
|
else:
|
|
if scale < 0.1:
|
|
raise ValueError("传入图像不能为空")
|
|
if height == 0:
|
|
height = 832
|
|
sub_height = height // 2
|
|
|
|
if scale > 0.01:
|
|
# 缩放并居中放置图像
|
|
def create_sub_image(height, scale, value):
|
|
H = int(height * scale)
|
|
# 计算目标位置
|
|
T = (height - H) // 2
|
|
# 创建目标大小的全零图像，形状为 [C, H, W]
|
|
result = torch.zeros((4, height, height), dtype=torch.uint8)
|
|
# 将缩放后的图像居中放置
|
|
result[:, T:T+H, T:T+H] = torch.full((4, H, H), value, dtype=torch.uint8)
|
|
return result
|
|
|
|
# 创建一个形状为 [C, H, W] 的组合图像
|
|
np_image = torch.zeros((4, sub_height, sub_height * 6), dtype=torch.uint8)
|
|
|
|
for i in range(6):
|
|
value = (i + 1) * 10
|
|
sub_image = create_sub_image(sub_height, scale, value)
|
|
# 将子图像拼接到 np_image 中
|
|
np_image[:, :, i * sub_height:(i + 1) * sub_height] = sub_image
|
|
equ = py360.cube2equi(
|
|
cubemap = np_image,
|
|
height = height,
|
|
width = height * 2,
|
|
mode = order,
|
|
cube_format='horizon')
|
|
|
|
equ_mask = equ.permute(1, 2, 0).unsqueeze(0)
|
|
|
|
# 转化输入为空时候，使用equ_mask返回
|
|
if np.size(result) == 0:
|
|
result = equ_mask / 255.0
|
|
result = result.to(torch.float32)
|
|
|
|
#提取透明通道转换为mask
|
|
mask = equ_mask[:, :, :, 3]
|
|
mask = (mask != 0).float()
|
|
|
|
#提取颜色通道转换为mask
|
|
masks = []
|
|
blue_channel = equ_mask[0][:, :, 2]
|
|
for i in range(6):
|
|
# 直接使用蓝色通道的值来创建遮罩
|
|
matching_pixels = (blue_channel == (i+1)*10).float()
|
|
|
|
# 转换为 PyTorch 张量
|
|
masks.append(matching_pixels)
|
|
|
|
masks = torch.stack(masks, dim=0)
|
|
|
|
if result.shape[-1] == 3:
|
|
result = add_alpha_channel(result)
|
|
|
|
if mask_split:
|
|
mask = result[:, :, :, 3]
|
|
mask_a = 1.0 - mask.squeeze(0)
|
|
mask_a = [mask_a] * 6
|
|
mask_a = torch.stack(mask_a, dim=0)
|
|
print(mask_a.shape, masks.shape)
|
|
masks = bitwise_and_masks(mask_a, masks)
|
|
|
|
masks = masks.unsqueeze(0)
|
|
|
|
return (result, mask, masks)
|
|
else:
|
|
# 当scale为0时，返回全透明
|
|
mask = result[:, :, :, 1]
|
|
masks = torch.zeros(1, 6, 64, 64)
|
|
return (result, mask, masks)
|
|
|
|
|
|
class PanoImageEqu2Cube:
|
|
def __init__(self):
|
|
pass
|
|
|
|
CATEGORY = "PanoCard/convert"
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"equ": ("IMAGE",),
|
|
"height": ("INT", {
|
|
"default": 832,
|
|
"min": 0,
|
|
"step": 8
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","MASK")
|
|
RETURN_NAMES = ("face","face_mask")
|
|
OUTPUT_IS_LIST = (True, True)
|
|
FUNCTION = "node"
|
|
|
|
def node(self, equ, height):
|
|
order = "nearest"
|
|
image = equ
|
|
sub_height, sub_width = image.shape[1], image.shape[2]
|
|
|
|
if sub_height*2 != sub_width:
|
|
raise ValueError("传入图像宽高必须为2:1")
|
|
|
|
if height == 0:
|
|
height = calculate_face_heght(sub_height)
|
|
|
|
Rot = [{'yaw': 0.0, 'pitch': 0.0, 'roll': 0.0}]
|
|
equ = image.permute(0, 3, 1, 2)
|
|
long_image = py360.equi2cube(
|
|
equi = equ,
|
|
rots = Rot,
|
|
w_face = height,
|
|
mode = order,
|
|
cube_format = 'horizon')
|
|
long_image = long_image.permute(0, 2, 3, 1)
|
|
|
|
# 分割图像
|
|
face_images = []
|
|
for i in range(6):
|
|
start_idx = i * height
|
|
end_idx = (i + 1) * height
|
|
face_image = long_image[0][:, start_idx:end_idx, :]
|
|
face_images.append(face_image)
|
|
|
|
# 将分割后的图像堆叠成一个新的张量
|
|
face_images = torch.stack(face_images, dim=0) # 形状变为 [6, H, H, C]
|
|
face = face_images.unsqueeze(0)
|
|
|
|
# 添加alpha通道
|
|
if face.shape[-1] == 3:
|
|
face = add_alpha_channel(face)
|
|
|
|
# 简单生成MASK
|
|
mask = face[:, :, :, :, -1]
|
|
|
|
return (face, mask)
|
|
|
|
def calculate_vfov(hfov, width, height):
|
|
# 将水平视场角从度转换为弧度
|
|
hfov_rad = hfov * np.pi / 180
|
|
|
|
# 计算水平方向的半视场角
|
|
half_hfov_rad = hfov_rad / 2
|
|
|
|
# 利用宽高比计算垂直方向的半视场角的正切值
|
|
tan_half_hfov = np.tan(half_hfov_rad)
|
|
tan_half_vfov = tan_half_hfov * (height / width)
|
|
|
|
# 计算垂直方向的半视场角
|
|
half_vfov_rad = np.arctan(tan_half_vfov)
|
|
|
|
# 计算垂直视场角并转换回度
|
|
vfov = 2 * half_vfov_rad * 180 / np.pi
|
|
|
|
return vfov
|
|
|
|
class PanoImageEqu2Equ:
|
|
def __init__(self):
|
|
pass
|
|
|
|
CATEGORY = "PanoCard/convert"
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"equ": ("IMAGE",),
|
|
"u_deg": ("FLOAT", {"default": 0.0,"min": -180,"max": 180,"step": 0.1}),
|
|
"v_deg": ("FLOAT", {"default": 0.0,"min": -90,"max": 90,"step": 0.1}),
|
|
"width": ("INT", {"default": 1664,"min": 0,"step": 8}),
|
|
"height": ("INT", {"default": 832,"min": 0,"step": 8}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
RETURN_NAMES = ("equ", )
|
|
FUNCTION = "node"
|
|
|
|
def node(self, equ, u_deg, v_deg, width, height):
|
|
image = equ.permute(0, 3, 1, 2)
|
|
|
|
if width == 0:
|
|
width = image.shape[3]
|
|
if height == 0:
|
|
height = image.shape[2]
|
|
|
|
# rotations
|
|
Rot = [{
|
|
'roll': 0.,
|
|
'pitch': np.pi * v_deg / 180, # rotate vertical
|
|
'yaw': np.pi * u_deg / 180, # rotate horizontal
|
|
}]
|
|
|
|
pic = py360.equi2equi(
|
|
src = image,
|
|
rots = Rot,
|
|
height = height,
|
|
width = width)
|
|
|
|
result = pic.permute(0, 2, 3, 1)
|
|
return (result,)
|
|
|
|
|
|
class PanoImagePic2Equ:
|
|
def __init__(self):
|
|
pass
|
|
|
|
CATEGORY = "PanoCard/convert"
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"hfov": ("FLOAT", {"default": 60.0, "min": 0.0, "max": 180.0, "step": 0.1}),
|
|
"u_deg": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
|
|
"v_deg": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1}),
|
|
"roll": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
|
|
"height": ("INT", {"default": 1024,"min": 8,"step": 8})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
RETURN_NAMES = ("equ", )
|
|
FUNCTION = "node"
|
|
|
|
def node(self, image, hfov, u_deg, v_deg, roll, height):
|
|
|
|
B = image.shape[0]
|
|
processed_images = []
|
|
|
|
x_yaw = -u_deg
|
|
y_pitch = -v_deg
|
|
|
|
for i in range(B):
|
|
# Process the image using the method
|
|
image_i = image[i]
|
|
vfov = calculate_vfov(hfov, image_i.shape[1], image_i.shape[0])
|
|
processed_image = plane_to_equirect(
|
|
image_i,
|
|
hfov,
|
|
vfov,
|
|
x_yaw,
|
|
y_pitch,
|
|
roll,
|
|
height + height
|
|
)
|
|
# Append the processed image to the list
|
|
processed_images.append(processed_image)
|
|
|
|
# Aggregate the processed images into a tensor of shape [B, H, W, C]
|
|
output_tensor = torch.stack(processed_images, dim=0)
|
|
|
|
return (output_tensor,)
|
|
|
|
class PanoImageEqu2Pic:
|
|
def __init__(self):
|
|
pass
|
|
|
|
CATEGORY = "PanoCard/convert"
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"equ": ("IMAGE",),
|
|
"u_deg": ("FLOAT", {"default": 0.0,"min": -180,"max": 180,"step": 0.1}),
|
|
"v_deg": ("FLOAT", {"default": 0.0,"min": -90,"max": 90,"step": 0.1}),
|
|
"hfov": ("FLOAT", {"default": 90.0,"min": 0,"max": 180,"step": 0.1}),
|
|
"width": ("INT", {"default": 512,"min": 8,"step": 8}),
|
|
"height": ("INT", {"default": 512,"min": 8,"step": 8}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
RETURN_NAMES = ("images", )
|
|
FUNCTION = "node"
|
|
|
|
def node(self, equ, u_deg, v_deg, hfov, width, height):
|
|
|
|
image = equ.permute(0, 3, 1, 2)
|
|
|
|
# rotations
|
|
Rot = [{
|
|
'roll': 0.,
|
|
'pitch': np.pi * v_deg / 180, # rotate vertical
|
|
'yaw': np.pi * u_deg / 180, # rotate horizontal
|
|
}]
|
|
|
|
pic = py360.equi2pers(
|
|
equi = image,
|
|
rots = Rot,
|
|
fov_x = hfov,
|
|
height = height,
|
|
width = width)
|
|
|
|
result = pic.permute(0, 2, 3, 1)
|
|
return (result,) |