# Path: ldif/util/file_util.py
# def readlines(p):
# def readbin(p):
# def writebin(p, s):
# def writetxt(p, s):
# def write_np(path, arr):
# def read_grd(path):
# def read_sif_v1(path, verbose=False):
# def read_lines(p):
# def read_image(p):
# def read_npz(p):
# def read_np(p):
# def read_txt_to_np(p):
# def read_py2_pkl(p):
# def write_mesh(path, mesh):
# def read_mesh(path):
# def read_csv(path):
# def read_normals(path_to_dir, im_count=20, leading='_depth'):
# def write_points(path_ext_optional, points):
# def write_grd(path, volume, world2grid=None):
# def write_depth_image(path, depth_image):
# 
# Path: ldif/util/geom_util_np.py
# def apply_4x4(arr, m, are_points=True, feature_count=0):
# def batch_apply_4x4(arrs, ms, are_points=True):
# def transform_normals(normals, tx):
# def world_xyzn_im_to_pts(world_xyz, world_n):
# def transform_r2n2_normal_cam_image_to_world_frame(normal_im, idx, e):
# def compute_argmax_image(xyz_image, decoder, embedding, k=1):
# def tile_world2local_frames(world2local, lyr):
# def extract_local_frame_images(world_xyz_im, world_normal_im, embedding,
#                                decoder):
# def select_top_k(argmax_im, xyzn_images):
# def make_r2n2_gt_top_k_image(e, idx, encoder, decoder, k=2):
# def depth_to_xyz_image(depth_image, intrinsics, extrinsics):
# def depth_to_point_cloud(depth_image, intrinsics, extrinsics):
# def depth_image_to_cam_image(depth_image,
#                              fx=585.0,
#                              fy=585.0,
#                              cx=255.5,
#                              cy=255.5):
# def depth_images_to_cam_images(depth_images,
#                                fx=585.0,
#                                fy=585.0,
#                                cx=255.5,
#                                cy=255.5):
# 
# Path: ldif/util/np_util.py
# def batch_np(arr, batch_size):
# def make_coordinate_grid(height, width, is_screen_space, is_homogeneous):
# def make_coordinate_grid_3d(length, height, width, is_screen_space,
#                             is_homogeneous):
# def filter_valid(mask, vals):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def make_mask(im, thresh=0.0):
# def make_pixel_mask(im):
# def thresh_and_radius_to_distance(radius, thresh):
# def plot_rbfs_at_thresh(centers, radii, thresh=0.5):
# def plot_rbfs(centers, radii, scale=10.0):
# def cube_and_render(volume, thresh):
# def sample_surface(quadrics, centers, radii, length, height, width,
#                    renormalize):
# def visualize_prediction(quadrics,
#                          centers,
#                          radii,
#                          renormalize,
#                          thresh=0.0,
#                          input_volumes=None):
# 
# Path: ldif/util/py_util.py
# def compose(*fs):
# def maybe(x, f):
# def py2_temporary_directory():
# def x11_server():
# def merge(x, y):
# def merge_into(x, ys):
# 
# Path: ldif/util/path_util.py
# def get_path_to_ldif_root():
# def get_path_to_ldif_parent():
# def package_to_abs(path):
# def gaps_path():
# def test_data_path():
# def util_test_data_path():
# def create_test_output_dir():
# 
# Path: ldif/util/file_util.py
# def readlines(p):
# def readbin(p):
# def writebin(p, s):
# def writetxt(p, s):
# def write_np(path, arr):
# def read_grd(path):
# def read_sif_v1(path, verbose=False):
# def read_lines(p):
# def read_image(p):
# def read_npz(p):
# def read_np(p):
# def read_txt_to_np(p):
# def read_py2_pkl(p):
# def write_mesh(path, mesh):
# def read_mesh(path):
# def read_csv(path):
# def read_normals(path_to_dir, im_count=20, leading='_depth'):
# def write_points(path_ext_optional, points):
# def write_grd(path, volume, world2grid=None):
# def write_depth_image(path, depth_image):


import math
import os
import subprocess as sp
import tempfile
import numpy as np
from ldif.util import file_util
from ldif.util import geom_util_np
from ldif.util import np_util
from ldif.util import py_util
from ldif.util import path_util
from ldif.util.file_util import log
  Args:
    depth_dir: Path to the root directory containing the scn2img output images.
    im_count: The number of images to read. Will read images with indices
      range(im_count).

  Returns:
    Numpy array with shape [im_count, height, width]. Dimensions determined from
      file.
  """
  depth_ims = []
  for i in range(im_count):
    path = depth_path_name(depth_dir, i)
    depth_ims.append(read_depth_im(path))
  depth_ims = np.stack(depth_ims)
  assert len(depth_ims.shape) == 3
  return depth_ims


def transform_r2n2_depth_image_to_gaps_frame(depth_im, idx, e):
  """Transforms a depth image predicted in the r2n2 space to the GAPS space."""
  # depth_im = self.r2n2_depth_images[idx, ...]
  # TODO(kgenova) Add support for predicted extrinsics.
  is_valid = depth_im != 0.0
  is_valid = np.reshape(is_valid, [224, 224])

  # Depth2cam:
  cam_im = gaps_depth_image_to_cam_image(depth_im, e.r2n2_xfov[idx])
  # Cam2world
  xyz_im = geom_util_np.apply_4x4(
      cam_im, e.r2n2_cam2world[idx, ...], are_points=True)

# Path: ldif/util/file_util.py
# def readlines(p):
# def readbin(p):
# def writebin(p, s):
# def writetxt(p, s):
# def write_np(path, arr):
# def read_grd(path):
# def read_sif_v1(path, verbose=False):
# def read_lines(p):
# def read_image(p):
# def read_npz(p):
# def read_np(p):
# def read_txt_to_np(p):
# def read_py2_pkl(p):
# def write_mesh(path, mesh):
# def read_mesh(path):
# def read_csv(path):
# def read_normals(path_to_dir, im_count=20, leading='_depth'):
# def write_points(path_ext_optional, points):
# def write_grd(path, volume, world2grid=None):
# def write_depth_image(path, depth_image):
# 
# Path: ldif/util/geom_util_np.py
# def apply_4x4(arr, m, are_points=True, feature_count=0):
# def batch_apply_4x4(arrs, ms, are_points=True):
# def transform_normals(normals, tx):
# def world_xyzn_im_to_pts(world_xyz, world_n):
# def transform_r2n2_normal_cam_image_to_world_frame(normal_im, idx, e):
# def compute_argmax_image(xyz_image, decoder, embedding, k=1):
# def tile_world2local_frames(world2local, lyr):
# def extract_local_frame_images(world_xyz_im, world_normal_im, embedding,
#                                decoder):
# def select_top_k(argmax_im, xyzn_images):
# def make_r2n2_gt_top_k_image(e, idx, encoder, decoder, k=2):
# def depth_to_xyz_image(depth_image, intrinsics, extrinsics):
# def depth_to_point_cloud(depth_image, intrinsics, extrinsics):
# def depth_image_to_cam_image(depth_image,
#                              fx=585.0,
#                              fy=585.0,
#                              cx=255.5,
#                              cy=255.5):
# def depth_images_to_cam_images(depth_images,
#                                fx=585.0,
#                                fy=585.0,
#                                cx=255.5,
#                                cy=255.5):
# 
# Path: ldif/util/np_util.py
# def batch_np(arr, batch_size):
# def make_coordinate_grid(height, width, is_screen_space, is_homogeneous):
# def make_coordinate_grid_3d(length, height, width, is_screen_space,
#                             is_homogeneous):
# def filter_valid(mask, vals):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def make_mask(im, thresh=0.0):
# def make_pixel_mask(im):
# def thresh_and_radius_to_distance(radius, thresh):
# def plot_rbfs_at_thresh(centers, radii, thresh=0.5):
# def plot_rbfs(centers, radii, scale=10.0):
# def cube_and_render(volume, thresh):
# def sample_surface(quadrics, centers, radii, length, height, width,
#                    renormalize):
# def visualize_prediction(quadrics,
#                          centers,
#                          radii,
#                          renormalize,
#                          thresh=0.0,
#                          input_volumes=None):
# 
# Path: ldif/util/py_util.py
# def compose(*fs):
# def maybe(x, f):
# def py2_temporary_directory():
# def x11_server():
# def merge(x, y):
# def merge_into(x, ys):
# 
# Path: ldif/util/path_util.py
# def get_path_to_ldif_root():
# def get_path_to_ldif_parent():
# def package_to_abs(path):
# def gaps_path():
# def test_data_path():
# def util_test_data_path():
# def create_test_output_dir():
# 
# Path: ldif/util/file_util.py
# def readlines(p):
# def readbin(p):
# def writebin(p, s):
# def writetxt(p, s):
# def write_np(path, arr):
# def read_grd(path):
# def read_sif_v1(path, verbose=False):
# def read_lines(p):
# def read_image(p):
# def read_npz(p):
# def read_np(p):
# def read_txt_to_np(p):
# def read_py2_pkl(p):
# def write_mesh(path, mesh):
# def read_mesh(path):
# def read_csv(path):
# def read_normals(path_to_dir, im_count=20, leading='_depth'):
# def write_points(path_ext_optional, points):
# def write_grd(path, volume, world2grid=None):
# def write_depth_image(path, depth_image):


# Path: ldif/util/gaps_util.py
import math
import os
import subprocess as sp
import tempfile
import numpy as np
from ldif.util import file_util
from ldif.util import geom_util_np
from ldif.util import np_util
from ldif.util import py_util
from ldif.util import path_util
from ldif.util.file_util import log

  Args:
    depth_dir: Path to the root directory containing the scn2img output images.
    im_count: The number of images to read. Will read images with indices
      range(im_count).

  Returns:
    Numpy array with shape [im_count, height, width]. Dimensions determined from
      file.
  """
  depth_ims = []
  for i in range(im_count):
    path = depth_path_name(depth_dir, i)
    depth_ims.append(read_depth_im(path))
  depth_ims = np.stack(depth_ims)
  assert len(depth_ims.shape) == 3
  return depth_ims


def transform_r2n2_depth_image_to_gaps_frame(depth_im, idx, e):
  """Transforms a depth image predicted in the r2n2 space to the GAPS space."""
  # depth_im = self.r2n2_depth_images[idx, ...]
  # TODO(kgenova) Add support for predicted extrinsics.
  is_valid = depth_im != 0.0
  is_valid = np.reshape(is_valid, [224, 224])

  # Depth2cam:
  cam_im = gaps_depth_image_to_cam_image(depth_im, e.r2n2_xfov[idx])
  # Cam2world
  xyz_im = geom_util_np.apply_4x4(
      cam_im, e.r2n2_cam2world[idx, ...], are_points=True)

  mask = np_util.make_mask(depth_im)