# Path: ibrnet/data_loaders/data_utils.py
# def rectify_inplane_rotation(src_pose, tar_pose, src_img, th=40):
#     relative = np.linalg.inv(tar_pose).dot(src_pose)
#     relative_rot = relative[:3, :3]
#     r = R.from_matrix(relative_rot)
#     euler = r.as_euler('zxy', degrees=True)
#     euler_z = euler[0]
#     if np.abs(euler_z) < th:
#         return src_pose, src_img
# 
#     R_rectify = R.from_euler('z', -euler_z, degrees=True).as_matrix()
#     src_R_rectified = src_pose[:3, :3].dot(R_rectify)
#     out_pose = np.eye(4)
#     out_pose[:3, :3] = src_R_rectified
#     out_pose[:3, 3:4] = src_pose[:3, 3:4]
#     h, w = src_img.shape[:2]
#     center = ((w - 1.) / 2., (h - 1.) / 2.)
#     M = cv2.getRotationMatrix2D(center, -euler_z, 1)
#     src_img = np.clip((255*src_img).astype(np.uint8), a_max=255, a_min=0)
#     rotated = cv2.warpAffine(src_img, M, (w, h), borderValue=(255, 255, 255), flags=cv2.INTER_LANCZOS4)
#     rotated = rotated.astype(np.float32) / 255.
#     return out_pose, rotated
# 
# def get_nearest_pose_ids(tar_pose, ref_poses, num_select, tar_id=-1, angular_dist_method='vector',
#                          scene_center=(0, 0, 0)):
#     '''
#     Args:
#         tar_pose: target pose [3, 3]
#         ref_poses: reference poses [N, 3, 3]
#         num_select: the number of nearest views to select
#     Returns: the selected indices
#     '''
#     num_cams = len(ref_poses)
#     num_select = min(num_select, num_cams-1)
#     batched_tar_pose = tar_pose[None, ...].repeat(num_cams, 0)
# 
#     if angular_dist_method == 'matrix':
#         dists = batched_angular_dist_rot_matrix(batched_tar_pose[:, :3, :3], ref_poses[:, :3, :3])
#     elif angular_dist_method == 'vector':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         scene_center = np.array(scene_center)[None, ...]
#         tar_vectors = tar_cam_locs - scene_center
#         ref_vectors = ref_cam_locs - scene_center
#         dists = angular_dist_between_2_vectors(tar_vectors, ref_vectors)
#     elif angular_dist_method == 'dist':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         dists = np.linalg.norm(tar_cam_locs - ref_cam_locs, axis=1)
#     else:
#         raise Exception('unknown angular distance calculation method!')
# 
#     if tar_id >= 0:
#         assert tar_id < num_cams
#         dists[tar_id] = 1e3  # make sure not to select the target id itself
# 
#     sorted_ids = np.argsort(dists)
#     selected_ids = sorted_ids[:num_select]
#     # print(angular_dists[selected_ids] * 180 / np.pi)
#     return selected_ids


import os
import numpy as np
import imageio
import torch
import sys
import json
from torch.utils.data import Dataset
from .data_utils import rectify_inplane_rotation, get_nearest_pose_ids
        else:
            id_render = -1
            subsample_factor = 1

        rgb = imageio.imread(rgb_file).astype(np.float32) / 255.
        rgb = rgb[..., [-1]] * rgb[..., :3] + 1 - rgb[..., [-1]]
        img_size = rgb.shape[:2]
        camera = np.concatenate((list(img_size), render_intrinsics.flatten(),
                                 render_pose.flatten())).astype(np.float32)

        nearest_pose_ids = get_nearest_pose_ids(render_pose,
                                                train_poses,
                                                int(self.num_source_views*subsample_factor),
                                                tar_id=id_render,
                                                angular_dist_method='vector')
        nearest_pose_ids = np.random.choice(nearest_pose_ids, self.num_source_views, replace=False)

        assert id_render not in nearest_pose_ids
        # occasionally include input image
        if np.random.choice([0, 1], p=[0.995, 0.005]) and self.mode == 'train':
            nearest_pose_ids[np.random.choice(len(nearest_pose_ids))] = id_render

        src_rgbs = []
        src_cameras = []
        for id in nearest_pose_ids:
            src_rgb = imageio.imread(train_rgb_files[id]).astype(np.float32) / 255.
            src_rgb = src_rgb[..., [-1]] * src_rgb[..., :3] + 1 - src_rgb[..., [-1]]
            train_pose = train_poses[id]
            train_intrinsics_ = train_intrinsics[id]
            if self.rectify_inplane_rotation:

# Path: ibrnet/data_loaders/data_utils.py
# def rectify_inplane_rotation(src_pose, tar_pose, src_img, th=40):
#     relative = np.linalg.inv(tar_pose).dot(src_pose)
#     relative_rot = relative[:3, :3]
#     r = R.from_matrix(relative_rot)
#     euler = r.as_euler('zxy', degrees=True)
#     euler_z = euler[0]
#     if np.abs(euler_z) < th:
#         return src_pose, src_img
# 
#     R_rectify = R.from_euler('z', -euler_z, degrees=True).as_matrix()
#     src_R_rectified = src_pose[:3, :3].dot(R_rectify)
#     out_pose = np.eye(4)
#     out_pose[:3, :3] = src_R_rectified
#     out_pose[:3, 3:4] = src_pose[:3, 3:4]
#     h, w = src_img.shape[:2]
#     center = ((w - 1.) / 2., (h - 1.) / 2.)
#     M = cv2.getRotationMatrix2D(center, -euler_z, 1)
#     src_img = np.clip((255*src_img).astype(np.uint8), a_max=255, a_min=0)
#     rotated = cv2.warpAffine(src_img, M, (w, h), borderValue=(255, 255, 255), flags=cv2.INTER_LANCZOS4)
#     rotated = rotated.astype(np.float32) / 255.
#     return out_pose, rotated
# 
# def get_nearest_pose_ids(tar_pose, ref_poses, num_select, tar_id=-1, angular_dist_method='vector',
#                          scene_center=(0, 0, 0)):
#     '''
#     Args:
#         tar_pose: target pose [3, 3]
#         ref_poses: reference poses [N, 3, 3]
#         num_select: the number of nearest views to select
#     Returns: the selected indices
#     '''
#     num_cams = len(ref_poses)
#     num_select = min(num_select, num_cams-1)
#     batched_tar_pose = tar_pose[None, ...].repeat(num_cams, 0)
# 
#     if angular_dist_method == 'matrix':
#         dists = batched_angular_dist_rot_matrix(batched_tar_pose[:, :3, :3], ref_poses[:, :3, :3])
#     elif angular_dist_method == 'vector':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         scene_center = np.array(scene_center)[None, ...]
#         tar_vectors = tar_cam_locs - scene_center
#         ref_vectors = ref_cam_locs - scene_center
#         dists = angular_dist_between_2_vectors(tar_vectors, ref_vectors)
#     elif angular_dist_method == 'dist':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         dists = np.linalg.norm(tar_cam_locs - ref_cam_locs, axis=1)
#     else:
#         raise Exception('unknown angular distance calculation method!')
# 
#     if tar_id >= 0:
#         assert tar_id < num_cams
#         dists[tar_id] = 1e3  # make sure not to select the target id itself
# 
#     sorted_ids = np.argsort(dists)
#     selected_ids = sorted_ids[:num_select]
#     # print(angular_dists[selected_ids] * 180 / np.pi)
#     return selected_ids


# Path: ibrnet/data_loaders/nerf_synthetic.py
import os
import numpy as np
import imageio
import torch
import sys
import json
from torch.utils.data import Dataset
from .data_utils import rectify_inplane_rotation, get_nearest_pose_ids

        else:
            id_render = -1
            subsample_factor = 1

        rgb = imageio.imread(rgb_file).astype(np.float32) / 255.
        rgb = rgb[..., [-1]] * rgb[..., :3] + 1 - rgb[..., [-1]]
        img_size = rgb.shape[:2]
        camera = np.concatenate((list(img_size), render_intrinsics.flatten(),
                                 render_pose.flatten())).astype(np.float32)

        nearest_pose_ids = get_nearest_pose_ids(render_pose,
                                                train_poses,
                                                int(self.num_source_views*subsample_factor),
                                                tar_id=id_render,
                                                angular_dist_method='vector')
        nearest_pose_ids = np.random.choice(nearest_pose_ids, self.num_source_views, replace=False)

        assert id_render not in nearest_pose_ids
        # occasionally include input image
        if np.random.choice([0, 1], p=[0.995, 0.005]) and self.mode == 'train':
            nearest_pose_ids[np.random.choice(len(nearest_pose_ids))] = id_render

        src_rgbs = []
        src_cameras = []
        for id in nearest_pose_ids:
            src_rgb = imageio.imread(train_rgb_files[id]).astype(np.float32) / 255.
            src_rgb = src_rgb[..., [-1]] * src_rgb[..., :3] + 1 - src_rgb[..., [-1]]
            train_pose = train_poses[id]
            train_intrinsics_ = train_intrinsics[id]
            if self.rectify_inplane_rotation:

                train_pose, src_rgb = rectify_inplane_rotation(train_pose, render_pose, src_rgb)