# Path: ibrnet/data_loaders/data_utils.py
# def quaternion_about_axis(angle, axis):
#     """Return quaternion for rotation about axis.
#     """
#     quaternion = np.zeros((4, ), dtype=np.float64)
#     quaternion[:3] = axis[:3]
#     qlen = vector_norm(quaternion)
#     if qlen > _EPS:
#         quaternion *= math.sin(angle/2.0) / qlen
#     quaternion[3] = math.cos(angle/2.0)
#     return quaternion
# 
# def quaternion_matrix(quaternion):
#     """Return homogeneous rotation matrix from quaternion.
#     """
#     q = np.array(quaternion[:4], dtype=np.float64, copy=True)
#     nq = np.dot(q, q)
#     if nq < _EPS:
#         return np.identity(4)
#     q *= math.sqrt(2.0 / nq)
#     q = np.outer(q, q)
#     return np.array((
#         (1.0-q[1, 1]-q[2, 2],     q[0, 1]-q[2, 3],     q[0, 2]+q[1, 3], 0.0),
#         (    q[0, 1]+q[2, 3], 1.0-q[0, 0]-q[2, 2],     q[1, 2]-q[0, 3], 0.0),
#         (    q[0, 2]-q[1, 3],     q[1, 2]+q[0, 3], 1.0-q[0, 0]-q[1, 1], 0.0),
#         (                0.0,                 0.0,                 0.0, 1.0)
#         ), dtype=np.float64)
# 
# def random_crop(rgb, camera, src_rgbs, src_cameras, size=(400, 600), center=None):
#     h, w = rgb.shape[:2]
#     out_h, out_w = size[0], size[1]
#     if out_w >= w or out_h >= h:
#         return rgb, camera, src_rgbs, src_cameras
# 
#     if center is not None:
#         center_h, center_w = center
#     else:
#         center_h = np.random.randint(low=out_h // 2 + 1, high=h - out_h // 2 - 1)
#         center_w = np.random.randint(low=out_w // 2 + 1, high=w - out_w // 2 - 1)
# 
#     rgb_out = rgb[center_h - out_h // 2:center_h + out_h // 2, center_w - out_w // 2:center_w + out_w // 2, :]
#     src_rgbs = np.array(src_rgbs)
#     src_rgbs = src_rgbs[:, center_h - out_h // 2:center_h + out_h // 2,
#                center_w - out_w // 2:center_w + out_w // 2, :]
#     camera[0] = out_h
#     camera[1] = out_w
#     camera[4] -= center_w - out_w // 2
#     camera[8] -= center_h - out_h // 2
#     src_cameras[:, 4] -= center_w - out_w // 2
#     src_cameras[:, 8] -= center_h - out_h // 2
#     src_cameras[:, 0] = out_h
#     src_cameras[:, 1] = out_w
#     return rgb_out, camera, src_rgbs, src_cameras
# 
# def random_flip(rgb, camera, src_rgbs, src_cameras):
#     h, w = rgb.shape[:2]
#     h_r, w_r = src_rgbs.shape[1:3]
#     rgb_out = np.flip(rgb, axis=1).copy()
#     src_rgbs = np.flip(src_rgbs, axis=-2).copy()
#     camera[2] *= -1
#     camera[4] = w - 1. - camera[4]
#     src_cameras[:, 2] *= -1
#     src_cameras[:, 4] = w_r - 1. - src_cameras[:, 4]
#     return rgb_out, camera, src_rgbs, src_cameras


import sys
import os
import numpy as np
import imageio
import torch
import json
from PIL import Image
from torch.utils.data import Dataset
from .data_utils import quaternion_about_axis, quaternion_matrix, random_crop, random_flip
        w_max, h_max = 0, 0
        for id in nearby_view_ids_selected:
            rgb_path = all_rgb_paths[id]
            ref_rgb = imageio.imread(rgb_path).astype(np.float32) / 255.
            h_in_view, w_in_view = all_img_sizes[id]
            h_img, w_img = ref_rgb.shape[:2]
            ref_rgbs.append(ref_rgb)
            ref_intrinsics = all_intrinsics[id]
            if h_in_view != h_img or w_in_view != w_img:
                ref_intrinsics[0] *= w_img / w_in_view
                ref_intrinsics[1] *= h_img / h_in_view
            ref_c2w = all_c2w_mats[id]
            ref_camera = np.concatenate([list(ref_rgb.shape[:2]), ref_intrinsics.flatten(), ref_c2w.flatten()])
            ref_cameras.append(ref_camera)
            h, w = ref_rgb.shape[:2]
            w_max = max(w, w_max)
            h_max = max(h, h_max)

        ref_rgbs_np = np.ones((len(ref_rgbs), h_max, w_max, 3), dtype=np.float32)
        for i, ref_rgb in enumerate(ref_rgbs):
            orig_h, orig_w = ref_rgb.shape[:2]
            h_start = int((h_max - orig_h) / 2.)
            w_start = int((w_max - orig_w) / 2.)
            ref_rgbs_np[i, h_start:h_start+orig_h, w_start:w_start+orig_w] = ref_rgb
            ref_cameras[i][4] += (w_max - orig_w) / 2.
            ref_cameras[i][8] += (h_max - orig_h) / 2.
            ref_cameras[i][0] = h_max
            ref_cameras[i][1] = w_max

        ref_cameras = np.array(ref_cameras)

# Path: ibrnet/data_loaders/data_utils.py
# def quaternion_about_axis(angle, axis):
#     """Return quaternion for rotation about axis.
#     """
#     quaternion = np.zeros((4, ), dtype=np.float64)
#     quaternion[:3] = axis[:3]
#     qlen = vector_norm(quaternion)
#     if qlen > _EPS:
#         quaternion *= math.sin(angle/2.0) / qlen
#     quaternion[3] = math.cos(angle/2.0)
#     return quaternion
# 
# def quaternion_matrix(quaternion):
#     """Return homogeneous rotation matrix from quaternion.
#     """
#     q = np.array(quaternion[:4], dtype=np.float64, copy=True)
#     nq = np.dot(q, q)
#     if nq < _EPS:
#         return np.identity(4)
#     q *= math.sqrt(2.0 / nq)
#     q = np.outer(q, q)
#     return np.array((
#         (1.0-q[1, 1]-q[2, 2],     q[0, 1]-q[2, 3],     q[0, 2]+q[1, 3], 0.0),
#         (    q[0, 1]+q[2, 3], 1.0-q[0, 0]-q[2, 2],     q[1, 2]-q[0, 3], 0.0),
#         (    q[0, 2]-q[1, 3],     q[1, 2]+q[0, 3], 1.0-q[0, 0]-q[1, 1], 0.0),
#         (                0.0,                 0.0,                 0.0, 1.0)
#         ), dtype=np.float64)
# 
# def random_crop(rgb, camera, src_rgbs, src_cameras, size=(400, 600), center=None):
#     h, w = rgb.shape[:2]
#     out_h, out_w = size[0], size[1]
#     if out_w >= w or out_h >= h:
#         return rgb, camera, src_rgbs, src_cameras
# 
#     if center is not None:
#         center_h, center_w = center
#     else:
#         center_h = np.random.randint(low=out_h // 2 + 1, high=h - out_h // 2 - 1)
#         center_w = np.random.randint(low=out_w // 2 + 1, high=w - out_w // 2 - 1)
# 
#     rgb_out = rgb[center_h - out_h // 2:center_h + out_h // 2, center_w - out_w // 2:center_w + out_w // 2, :]
#     src_rgbs = np.array(src_rgbs)
#     src_rgbs = src_rgbs[:, center_h - out_h // 2:center_h + out_h // 2,
#                center_w - out_w // 2:center_w + out_w // 2, :]
#     camera[0] = out_h
#     camera[1] = out_w
#     camera[4] -= center_w - out_w // 2
#     camera[8] -= center_h - out_h // 2
#     src_cameras[:, 4] -= center_w - out_w // 2
#     src_cameras[:, 8] -= center_h - out_h // 2
#     src_cameras[:, 0] = out_h
#     src_cameras[:, 1] = out_w
#     return rgb_out, camera, src_rgbs, src_cameras
# 
# def random_flip(rgb, camera, src_rgbs, src_cameras):
#     h, w = rgb.shape[:2]
#     h_r, w_r = src_rgbs.shape[1:3]
#     rgb_out = np.flip(rgb, axis=1).copy()
#     src_rgbs = np.flip(src_rgbs, axis=-2).copy()
#     camera[2] *= -1
#     camera[4] = w - 1. - camera[4]
#     src_cameras[:, 2] *= -1
#     src_cameras[:, 4] = w_r - 1. - src_cameras[:, 4]
#     return rgb_out, camera, src_rgbs, src_cameras


# Path: ibrnet/data_loaders/spaces_dataset.py
import sys
import os
import numpy as np
import imageio
import torch
import json
from PIL import Image
from torch.utils.data import Dataset
from .data_utils import quaternion_about_axis, quaternion_matrix, random_crop, random_flip

        w_max, h_max = 0, 0
        for id in nearby_view_ids_selected:
            rgb_path = all_rgb_paths[id]
            ref_rgb = imageio.imread(rgb_path).astype(np.float32) / 255.
            h_in_view, w_in_view = all_img_sizes[id]
            h_img, w_img = ref_rgb.shape[:2]
            ref_rgbs.append(ref_rgb)
            ref_intrinsics = all_intrinsics[id]
            if h_in_view != h_img or w_in_view != w_img:
                ref_intrinsics[0] *= w_img / w_in_view
                ref_intrinsics[1] *= h_img / h_in_view
            ref_c2w = all_c2w_mats[id]
            ref_camera = np.concatenate([list(ref_rgb.shape[:2]), ref_intrinsics.flatten(), ref_c2w.flatten()])
            ref_cameras.append(ref_camera)
            h, w = ref_rgb.shape[:2]
            w_max = max(w, w_max)
            h_max = max(h, h_max)

        ref_rgbs_np = np.ones((len(ref_rgbs), h_max, w_max, 3), dtype=np.float32)
        for i, ref_rgb in enumerate(ref_rgbs):
            orig_h, orig_w = ref_rgb.shape[:2]
            h_start = int((h_max - orig_h) / 2.)
            w_start = int((w_max - orig_w) / 2.)
            ref_rgbs_np[i, h_start:h_start+orig_h, w_start:w_start+orig_w] = ref_rgb
            ref_cameras[i][4] += (w_max - orig_w) / 2.
            ref_cameras[i][8] += (h_max - orig_h) / 2.
            ref_cameras[i][0] = h_max
            ref_cameras[i][1] = w_max

        ref_cameras = np.array(ref_cameras)

        if self.mode == 'train' and self.random_crop: