launch_compare_random_tg.py

'''
Randomly choose an image(here image 000063) in our dataset as the default-pose-image in sds.
Run sds with the chosen image.
Randomly choose a camera view in the 240 pics generated by sds as a target view.
Use the Matrix_nerf_to_dataset(M_n2d) to convert target view pose to dataset coordinate system.
Choose the nearest camera view in dataset and generate novel view.

Note: still cannot align nerf's coordinate and dataset's coordinate.
ref: data/rand_tg_compare_output/6_test_nd_convert_datasetcoord/view_camera_y-x.png
'''

import os
from PIL import Image

import cv2
import numpy as np
import torch
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.cm as cm

from ldm.models.diffusion import options
options.LDM_DISTILLATION_ONLY = True

from threestudio.models.guidance import zero123_guidance
from omegaconf import OmegaConf 

from test_camera import test_camera, get_defaults
from view_camera import plot_camera_pose


# def compute_m_d2n(nerf_pose, dataset_pose):
#     '''
#     Matrix, dataset coordinate to nerf coordinate
#     dataset=A, nerf=B
#     T_B = T * T_A 
#     T = T_B * T_A_inv
#     '''
#     # Compute T_A_inv
#     dataset_r = dataset_pose[:3, :3]
#     dataset_t = dataset_pose[:3, 3]
#     dataset_r_inv = dataset_r.T
#     dataset_t_inv = -dataset_r_inv @ dataset_t
#     dataset_pose_inv = np.eye(4)
#     dataset_pose_inv[:3, :3] = dataset_r_inv
#     dataset_pose_inv[:3, 3] = dataset_t_inv
    
#     M_d2n = nerf_pose @ dataset_pose_inv
#     return M_d2n


# def compute_m_n2d(dataset_pose, nerf_pose):
#     '''
#     Matrix dataset coordinate to dataset coordinate
#     '''
#     nerf_r = nerf_pose[:3, :3]
#     nerf_t = nerf_pose[:3, 3]
#     nerf_r_inv = nerf_r.T
#     nerf_t_inv = -nerf_r_inv @ nerf_t
#     nerf_pose_inv = np.eye(4)
#     nerf_pose_inv[:3, :3] = nerf_r_inv
#     nerf_pose_inv[:3, 3] = nerf_t_inv
    
#     M_n2d = dataset_pose @ nerf_pose_inv
#     return M_n2d


def compute_m_d2n(nerf_pose, dataset_pose):
    dataset_pose_inv = np.linalg.inv(dataset_pose)
    M_d2n = nerf_pose @ dataset_pose_inv
    return M_d2n

def compute_m_n2d(dataset_pose, nerf_pose):
    nerf_pose_inv = np.linalg.inv(nerf_pose)
    M_n2d = dataset_pose @ nerf_pose_inv
    return M_n2d


# 'gl9' inverse camera's z axis
# CORRECT, according to view_camera
# however sometimes z will be error
def opencv_to_opengl(pose):
    # 定义反射矩阵 S
    S = np.diag([1, 1, -1])

    R = pose[:3,:3]

    # 计算新的旋转矩阵 R'
    R_prime = R @ S

    pose_gl = pose.copy()
    pose_gl[:3,:3] = R_prime
    return pose_gl


def load_K_Rt_from_P(filename, P=None):
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]
    euler_angle = out[6]

    K = K/K[2,2]
    intrinsics = np.eye(4) 
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3,3] = (t[:3] / t[3])[:,0] 

    return intrinsics, pose, euler_angle


# OpenCV
def get_camera_poses(cam_file_path):
    camera_dict = np.load(cam_file_path)
    n_images = len(camera_dict.files) // 2
    scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
    world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]

    intrinsics_all = []
    pose_all = []
    # forward_vectors_world_all = []  # 用于存储世界坐标系下的前向向量
    camera_positions_world_all = []  # 用于存储世界坐标系下的相机位置
    euler_angles_all = []  # 用于存储前向向量的欧拉角
    euler_angles_returned_all = []
    save_data = {}
    idx = 0
    for scale_mat, world_mat in zip(scale_mats, world_mats):
        P = world_mat @ scale_mat 
        P = P[:3, :4]
        intrinsics, pose, euler_angle_returned = load_K_Rt_from_P(None, P)

        pose = opencv_to_opengl(pose) 

        # because we do resize and center crop 384x384 when using omnidata model, we need to adjust the camera intrinsic accordingly
        scale = 384 / 680 # raw is 680 and resized image is 384
        offset = (1200 - 680 ) * 0.5 # photocentre, width from 1200 to 680
        intrinsics[0, 2] -= offset
        intrinsics[:2, :] *= scale

        intrinsics_all.append(intrinsics)
        pose_all.append(pose)

        # 将内参和位姿矩阵分别以 'intrinsics_x' 和 'pose_x' 命名保存
        save_data[f'intrinsics_{idx}'] = intrinsics
        save_data[f'pose_{idx}'] = pose


        # 计算相机前向向量在世界坐标系下的方向
        # cam_to_world = np.eye(4)  # 扩展 3x4 矩阵为 4x4 齐次坐标矩阵
        # # cam_to_world[:3, :4] = pose
        cam_to_world = pose

        # 相机前向向量 (0, 0, 1) 在相机坐标系下
        # forward_vector_cam = np.array([0, 0, -1, 0])  # 使用齐次坐标
        forward_vector_cam = np.array([0, 0, 1, 0])  # 使用齐次坐标
        up_vector_cam = np.array([0, 1, 0, 0])  # 使用齐次坐标

        # # 通过位姿矩阵将前向向量转换到世界坐标系
        forward_vector_world = cam_to_world @ forward_vector_cam
        up_vector_world = cam_to_world @ up_vector_cam
        # forward_vectors_world_all.append(forward_vector_world[:3])  # 存储前向向量

        # # 保存前向向量
        # save_data[f'forward_vector_{idx}'] = forward_vector_world[:3]  # 保存转换后的前向向量


        # # 计算前向向量的欧拉角，并保存
        # roll_deg, pitch_deg, yaw_deg = vector_to_euler(forward_vector_world[:3], up_vector_world[:3]) # xyz
        # euler_angles_all.append([roll_deg, pitch_deg, yaw_deg])
        # # 保存欧拉角
        # save_data[f'forward_vector_euler_compute{idx}'] = [roll_deg, pitch_deg, yaw_deg]


        # 保存opencv decompose 返回的欧拉角
        euler_angles_returned_all.append(euler_angle_returned)
        save_data[f'forward_vector_euler_returned{idx}'] = euler_angle_returned



        # 相机位置向量 (0, 0, 0) 在相机坐标系下，用齐次坐标表示
        position_cam = np.array([0, 0, 0, 1])  # 使用齐次坐标，最后一位是 1 表示位置向量

        # 通过位姿矩阵将相机位置转换到世界坐标系
        position_world = pose @ position_cam
        camera_positions_world_all.append(position_world[:3])  # 存储相机位置向量

        # 保存相机位置
        save_data[f'camera_position_{idx}'] = position_world[:3]  # 保存转换后的相机位置
        idx += 1


    save_path = os.path.join("data/image_output", "20_poses_position_euler2.npz")
    np.savez(save_path, **save_data)
    # np.savez(save_path, intrinsics_all=intrinsics_all, pose_all=pose_all)

    print('camera params loaded') 
    return intrinsics_all, pose_all, euler_angles_returned_all


def get_cond_camera_poses(pose_all):
    # all images are condition
    cond_poses = []
    cond_idx = []
    input_camera_spacing = 1                                  
    for idx, pose in enumerate(pose_all):
        if idx % input_camera_spacing == 0:
            cond_poses.append(pose)
            cond_idx.append(idx)
    return cond_poses, cond_idx
    

def get_target_camera_poses(pose_all):
    target_poses = []
    target_idx = []
    input_camera_spacing = 6                                
    for idx, pose in enumerate(pose_all):
        if idx % input_camera_spacing != 0:
            target_poses.append(pose)
            target_idx.append(idx)
    return target_poses, target_idx


def extract_pose_components(pose):
    translation = pose[:3, 3]
    rotation = pose[:3, :3]
    return translation, rotation


def get_forward_vector_from_pose(pose):
    direction = -pose[:3, 2]  # forward_vector
    return direction

def get_forward_vector_from_pose_2d(pose):
    direction = -pose[:3, 2]  # forward_vector # array
    direction2d = direction[1:3]
    return direction2d


def angle_between_vectors(u, v):
    u = np.array(u)
    v = np.array(v)
    dot_product = np.dot(u, v)
    norm_u = np.linalg.norm(u)
    norm_v = np.linalg.norm(v)
    if norm_u == 0 or norm_v == 0:
        raise ValueError("零向量不能计算夹角。")
    cos_theta = dot_product / (norm_u * norm_v)
    cos_theta = np.clip(cos_theta, -1.0, 1.0)  # 防止数值误差
    theta_rad = np.arccos(cos_theta)
    theta_deg = np.degrees(theta_rad)
    return theta_deg


def calculate_distance(input_pose, candidate_pose, alpha=1.0, beta=1.0):
    input_translation, input_rotation = extract_pose_components(input_pose)
    candidate_translation, candidate_rotation = extract_pose_components(candidate_pose)
    input_fv = get_forward_vector_from_pose_2d(input_pose)
    candidate_fv = get_forward_vector_from_pose_2d(candidate_pose)
    theta = angle_between_vectors(input_fv, candidate_fv)
    
    position_distance = np.linalg.norm(input_translation - candidate_translation)
    
    return alpha * position_distance + beta * np.abs(theta) 


def find_nearest_cond(target_pose, cond_poses, cond_idx, alpha=1.0, beta=1.0, target_euler=None, euler_all=None): # cond euler
    '''
    cond poses: a list
    cond idx: a list, the raw idx of each cond pose
    return the nearest cond pose of target pose, and the raw idx of it
    '''
    nearest_cond_pose = None
    nearest_cond_idx = None
    min_distance = float('inf')

    for pose, idx in zip(cond_poses, cond_idx):
        distance = calculate_distance(target_pose, pose, alpha, beta)
        if distance < min_distance:
            min_distance = distance
            nearest_cond_pose = pose
            nearest_cond_idx = idx

    return nearest_cond_pose, nearest_cond_idx


def magnify_dataset_translation(poses):
    '''
    to align with nerf's translation.
    input a list.
    '''
    magnified_all = []
    for i, p in enumerate(poses):
        magnified = p
        magnified[:3, 3] = p[:3, 3] * 10
        magnified_all.append(magnified)
    return magnified_all
    # return poses


def launch():

    default_img_path = 'data/image_test/000063_rgb.png'
    img_path = 'data/image_test'
    sds_img_path = 'data/sds_outputs/save2/it4000-val'
    output_path = 'data/rand_tg_compare_output/6_test_nd_convert_datasetcoord'
    translation_weight = 0
    rotation_weight = 1

    # get cond poses
    cam_file = 'data/cameras.npz'
    intrinsic_all, pose_all, euler_returned_all = get_camera_poses(cam_file) # already gl dataset pose.
    cond_poses, cond_idx = get_cond_camera_poses(pose_all) # all images are condition. cond_poses equals to pose_all.
    cond_poses_mag = magnify_dataset_translation(cond_poses)
    
    # cond_euler = [euler_returned_all[i] for i in cond_idx]

    # get fov
    intrinsic = intrinsic_all[0]
    f_x = intrinsic[0, 0] 
    image_width = 256
    fov_horizontal = 2 * np.arctan(image_width / (2 * f_x)) * 180 / np.pi  # rad -> degree
    fov_tensor = torch.from_numpy(np.array([fov_horizontal])).cuda().to(torch.float32)

    # align nerf coord and dataset coord, and get target poses
    # target_poses, target_idx = get_target_camera_poses(pose_all) # when targets are chosen from dataset. here we don't use this function.
    target_space = 10
    target_poses_nerf_all = test_camera()
    target_poses_nerf_all = np.array(target_poses_nerf_all)
    target_poses_nerf = [target_poses_nerf_all[i] for i in range(len(target_poses_nerf_all)) if i % target_space == 0]
    default_dataset_pose = cond_poses_mag[63]
    default_nerf_pose = target_poses_nerf_all[0]
    m_n2d = compute_m_n2d(default_dataset_pose, default_nerf_pose)
    target_poses = [m_n2d @ target_poses_nerf[i] for i in range(len(target_poses_nerf))]


    if not os.path.exists(output_path):
        os.makedirs(output_path)


    # view camera
    # using sds's pose can generate correct view, which suggests the figure's coordinate system is aligned with the sds's coordinate system.
    view_camera_save_path_y_x = os.path.join(output_path, "view_camera_y-x.png")
    view_camera_save_path_z_y = os.path.join(output_path, "view_camera_z-y.png")
    view_space = 1 # our_dataset:1  sds:10
    poses = cond_poses_mag
    # poses.append(target_poses) # this will cause poses's last element becomes a tuple.
    poses.extend(target_poses)
    our_dataset = True # a scale in translation

    # cond_poses_mag = magnify_dataset_translation(cond_poses)
    # default_dataset_pose = cond_poses_mag[63]
    # m_d2n = compute_m_d2n(default_nerf_pose, default_dataset_pose)
    # cond_poses_nerf = [m_d2n @ cond_poses_mag[i] for i in range(len(cond_poses_mag))]
    # poses = target_poses_nerf
    # poses.extend(cond_poses_nerf)

    num_cameras = len(poses)
    cmap = cm.get_cmap('Spectral')  # colormap
    colors = [cmap(i / num_cameras) for i in range(num_cameras)] 
        
    arrow_length = 5
    fig, ax = plt.subplots(figsize=(10, 8))
    for i, T in enumerate(poses):
        if i % view_space == 0:
            label = f'C{i}'
            color = colors[i % len(colors)]  
            plot_camera_pose(T, ax=ax, label=label, color=color, our_dataset=our_dataset, vert_ax='y', hori_ax='x', arrow_length=arrow_length)
    plt.savefig(view_camera_save_path_y_x)
    plt.clf()

    fig, ax = plt.subplots(figsize=(10, 8))
    for i, T in enumerate(poses):
        if i % view_space == 0:
            label = f'C{i}'
            color = colors[i % len(colors)]  
            plot_camera_pose(T, ax=ax, label=label, color=color, our_dataset=our_dataset, vert_ax='z', hori_ax='y', arrow_length=arrow_length)
    plt.savefig(view_camera_save_path_z_y)
    print("view camera figure saved")


    for i, tp in enumerate(target_poses):
        target_pose = tp
        idx = i * target_space
        # target_euler = euler_returned_all[idx]
        nearest_cond_pose, nearest_cond_idx = find_nearest_cond(target_pose, cond_poses, cond_idx, alpha=translation_weight, beta=rotation_weight) # XXX

        target_pose_gl = target_pose
        nearest_cond_pose_gl = nearest_cond_pose

        cond_img_path = os.path.join(img_path, f'{nearest_cond_idx:06d}_rgb.png')

        guidance_cfg = dict(
            pretrained_model_name_or_path= "zeronvs.ckpt", 
            pretrained_config= "zeronvs_config.yaml", 
            guidance_scale= 7.5,
            cond_image_path =cond_img_path,
            min_step_percent=[0,.75,.02,1000],
            max_step_percent=[1000, 0.98, 0.025, 2500],
            vram_O=False 
        )
        guidance = zero123_guidance.Zero123Guidance(OmegaConf.create(guidance_cfg))

        cond_image_pil = Image.open(cond_img_path).convert("RGB")
        cond_image_pil = cond_image_pil.resize((256, 256)) # XXX
        cond_image = torch.from_numpy(np.array(cond_image_pil)).cuda() / 255.

        c_crossattn, c_concat = guidance.get_img_embeds(
            cond_image.permute((2, 0, 1))[None]) # change (H, W, C) to (C, H, W)
        
        cond_camera = nearest_cond_pose_gl
        target_camera = target_pose_gl
        cond_camera = torch.from_numpy(cond_camera[None]).cuda().to(torch.float32)
        target_camera = torch.from_numpy(target_camera[None]).cuda().to(torch.float32)
        camera_batch = {
            "target_cam2world": target_camera, 
            "cond_cam2world": cond_camera,
            # "fov_deg": torch.from_numpy(np.array([45.0])).cuda().to(torch.float32) 
            "fov_deg": fov_tensor # XXX horizontal field of view, degree
        }

        guidance.cfg.precomputed_scale=.7
        cond = guidance.get_cond_from_known_camera(
            camera_batch,
            c_crossattn=c_crossattn,
            c_concat=c_concat,
            # precomputed_scale=.7,
        )

        print("------camerabatch--------")
        print(camera_batch["cond_cam2world"]) 
        print(camera_batch["target_cam2world"]) 

        novel_view = guidance.gen_from_cond(cond) 
        novel_view_pil = Image.fromarray(np.clip(novel_view[0]*255, 0, 255).astype(np.uint8))

        # target image gt(take sds as gt)
        sds_novel_path = os.path.join(sds_img_path, f'{i}.png')
        sds_image_pil = Image.open(sds_novel_path).convert("RGB")
        width, height = sds_image_pil.size
        split_width = width // 4
        split_sds_image = sds_image_pil.crop((0, 0, split_width, height)).resize((256, 256)) # (left, upper, right, lower)

        default_pil = Image.open(default_img_path).convert("RGB").resize((256, 256))

        default_image_array = np.array(default_pil)
        sds_image_array = np.array(split_sds_image)
        cond_image_array = np.array(cond_image_pil)
        novel_image_array = np.array(novel_view_pil)


        if not os.path.exists(os.path.join(output_path, 'concat')):
            os.makedirs(os.path.join(output_path, 'concat'))
        concatenated = np.hstack((default_image_array, sds_image_array, cond_image_array, novel_image_array))
        print("concat image:", concatenated)
        concatenated_image = Image.fromarray(concatenated)
        concatenated_save_path = os.path.join(output_path, f'concat/{i}_{idx}tg_{nearest_cond_idx}c.png')
        print("concat save path: ", concatenated_save_path)
        concatenated_image.save(concatenated_save_path)


        


if __name__ == '__main__':
    launch()