def _make_optimized_dataset(directory, batch_size, mode, split):
def build_dataset_obj(dataset_items, bs):
def make_dataset(directory, batch_size, mode, split): def _setup_cam(camera):
def ptsview(pts, mesh=None, camera='fixed'):
def mshview(mesh, camera='fixed'):
def grdview(volume, world2grid=None):
def gapsview(keep=True, **kwargs):
def _make_pts_input_str(d, pts, allow_none=True):
def _gapsview(d,
              msh=None,
              pts=None,
              grd=None,
              world2grid=None,
              grid_threshold=0.0,
              camera='default'):
def read_pts_file(path):
def read_cam_file(path, verbose=False):
def batch_gaps_depth_image_to_cam_image(depth_image, xfov=None):
def gaps_depth_image_to_cam_image(depth_image, xfov=0.5):
def read_depth_im(path):
def read_normals_im(px, py, pz):
def read_normals_dir(directory, image_count):
def depth_path_name(depth_dir, idx):
def read_depth_directory(depth_dir, im_count):
def transform_r2n2_depth_image_to_gaps_frame(depth_im, idx, e): def _parse_synset_or_cat(synset_or_cat):
  def __init__(self, is_local):
  def random_example(self, split, synset_or_cat=None, equal_class_freq=True):
def assert_is_4x4(m):
def _get_world_pts_from_idx(idx,
                            depth,
                            world_xyz_im,
                            world_n_im,
                            target_point_count=10000):
  def __init__(self, split, synset_or_cat, mesh_hash=None, dynamic=False,
               verbose=True):
  def from_relpath(cls, relpath):
  def from_local_dataset_tokens(cls, model_directory, split, class_name,
                                mesh_name):
  def from_directory(cls, dirpath, verbose=True):
  def from_npz_path(cls, npz_path):
  def from_elt(cls, elt):
  def npz_path(self):
  def surface_samples_path(self):
  def grid_path(self):
  def dodeca_depth_and_normal_path(self):
  def r2n2_depth_dir(self):
  def r2n2_cam2v1_path(self):
  def v12occnet_path(self):
  def from_rgb_path_and_split(cls, rgb_path, split):
  def set_rgb_path(self, rgb_path):
  def set_rgb_path_from_idx(self, idx):
  def _archive(self):
  def rgb_path(self):
  def set_depth_res(self, res):
  def rgb_image(self):
  def r2n2_images(self):
  def load_depth_and_normal_npz(self, path):
  def r2n2_depth_images(self):
  def r2n2_normal_world_images(self):
  def r2n2_normal_cam_images(self):
  def r2n2_cam_images(self):
  def r2n2_xyz_images(self):
  def set_r2n2_cam2v1(self, cam2v1):
  def r2n2_cam2v1(self):
  def r2n2_xfov(self):
  def r2n2_viewpoints(self):
      def compute_pose(metadata_line):
  def v1_watertight_gt_mesh(self):
  def v12occnet(self):
  def occnet2v1(self):
  def gaps2occnet(self):
  def gaps2v1(self):
  def depr_depth_images(self):
  def depth_images(self):
  def cam_normal_images(self):
  def world_normal_images(self):
  def surface_samples(self):
  def dodeca_cam2world(self):
  def cam_xyz_images_from_dodeca(self):
  def world_xyz_images_from_dodeca(self):
  def get_surface_samples_from_single_dodeca_image(self, idx):
  def get_surface_samples_from_single_r2n2_depth_image(self, idx):
  def surface_sample_count(self):
  def surface_sample_count(self, count):
  def precomputed_surface_samples_from_dodeca(self):
  def surface_samples_from_dodeca(self):
  def get_surface_samples(self, sample_count):
  def near_surface_samples(self):
  def uniform_samples(self):
  def mesh_name(self):
  def grid(self):
  def world2grid(self):
  def gt_path(self):
  def normalized_gt_mesh(self):
  def v1_gt_mesh(self):
  def tx_path(self):
  def r2n2_cam2world(self):
  def occnet2gaps(self):
  def tx(self):
  def gaps_to_occnet(self):
  def gt_mesh(self):
  def gt_mesh_str(self):
  def max_depth_normals_path(self):
  def _load_max_depth_normals(self):
  def max_depth_512(self):
  def max_pred_cam_normals_512(self):
  def max_cam_xyz(self):
  def bts_cam_xyz(self):
  def bts_cam2world(self):
  def max_cam2occnet_path(self):
  def max_cam2occnet(self):
  def max_cam2world(self):
  def max_world_xyz(self):
  def bts_world_xyz(self):
  def max_world_normals(self):
  def bts_world_normals(self):
  def get_bts_world_pts_from_idx(self, idx):
  def get_max_world_pts_from_idx(self, idx):
  def _batch_resize(self, ims, res, strategy='nearest'):
  def max_world_xyz_224(self):
  def bts_world_xyz_224(self):
  def max_world_normals_224(self):
  def bts_depth_224(self):
  def bts_depth_480(self):
  def bts_pred_cam_normals_480(self):
  def max_depth_224(self):
class ExampleGenerator(object):
class InferenceExample(object): def backwards_compatible_hparam_defaults(ident):
def tf_hparams_to_dict(tf_hparams):
def build_ldif_hparams():
def build_sif_hparams():
def build_improved_sif_hparams():
def build_singleview_depth_hparams():
def write_hparams(hparams, path):
def read_hparams_with_new_backwards_compatible_additions(path):
def read_hparams(path): def to_xid(job):
  def __init__(self, parent_job, idx):
  def relpath(self):
  def directory(self):
  def abspath(self):
  def __init__(self, parent_experiment, job_str, use_temp_ckpts='warn'):
  def from_xid(cls, parent_experiment, xid):
  def _ensure_lazy_attrs_are_zero(self):
  def set_use_temp_ckpts(self, useable):
  def is_visible(self):
  def ensure_visible(self):
  def ckpt_dir(self):
  def root_dir(self):
  def all_checkpoint_indices(self):
  def all_checkpoints(self):
  def latest_checkpoint_before(self, idx, must_equal=False):
  def has_at_least_one_checkpoint(self):
  def ensure_has_at_least_one_checkpoint(self):
  def checkpoint_count(self):
  def newest_checkpoint(self):
  def k_evenly_spaced_checkpoints(self, k):
  def hparams(self):
  def model_config(self):
  def __init__(self, model_dir, model_name, experiment_name):
  def set_order(self, ordered_xids):
  def job_with_xid(self, xid):
  def filter_jobs_by_xid(self, xids):
  def visible_jobs(self):
  def visible_job_count(self):
  def visible_xids(self):
  def job_from_xmanager_id(self, xid, must_be_visible=True):
  def job_from_job_str(self, job_str, must_be_visible=True):
  def __init__(self, experiment, desired_ckpts=-1):
  def local_root(self):
  def remote_root(self):
  def remote_result_ckpt_dir(self, xid):
  def local_result_ckpt_dir(self, xid):
  def ckpts_for_xid_and_split(self, xid, split):
  def ckpt_for_xid(self, xid, split):
  def remote_result_base(self, xid, split):
  def all_remote_mesh_names(self, split, ensure_nonempty=True):
  def copy_meshes(self, mesh_names, split, overwrite_if_present=True):
  def random_remote_mesh_names(self, count, split):
  def remote_path_to_mesh(self, mesh_name, xid, split):
  def local_path_to_mesh(self, mesh_name, xid, split):
  def _mesh_relative_path(self, mesh_name):
  def local_qual_path(self):
  def __init__(self, hparams):
class Checkpoint(object):
class Job(object):
class Experiment(object):
class ResultStore(object):
class ModelConfig(object): FLAGS = flags.FLAGS
def set_train_op(model_config):
def sif_transcoder(model_config): 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): def get_free_gpu_memory(cuda_device_index):
def get_allowable_fraction_without(mem_to_reserve, cuda_device_index): 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(): 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 os
import sys
import time
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from ldif.datasets import local_inputs
from ldif.util import gaps_util
from ldif.inference import example
from ldif.model import hparams
from ldif.inference import experiment as experiments
from ldif.training import shared_launcher
from ldif.util import file_util
from ldif.util import gpu_util
from ldif.util import path_util
from ldif.util.file_util import log
flags.DEFINE_string(
    'experiment_name', 'reproduce-ldif',
    'The name of the experiment. It will be saved to disk under this name,'
    ' and can be reloaded for inference, eval, interactive notebooks, or'
    ' further training by providing this experiment_name. To resume training'
    ' from an already partially trained model, simply provide an experiment'
    ' name that is already written. Note that global step will be reloaded.')

flags.DEFINE_integer(
    'train_step_count', 1000000,
    'The number of training steps to take before training is complete.')

flags.DEFINE_string(
    'model_directory', 'trained_models/',
    'The path to the trained model root directory. Can be'
    ' absolute or relative to the LDIF repository root.')

flags.DEFINE_string(
    'dataset_directory', '',
    'The path to the dataset generated by meshes2dataset.py'
    ' Must be provided. The train split will be used.')

flags.DEFINE_string('split', 'train', 'The split to train on.')

flags.DEFINE_boolean(
    'visualize', False, 'If true, visualizes the training'
    ' data before training. Useful for debugging data errors.')

flags.DEFINE_string('model_type', 'ldif', 'The type of model to train. Can be'
                    'ldif (the LDIF paper model), sif (the SIF paper)'
                    ', or sif++ (an improved version of SIF using the LDIF'
                    ' architecture changes)')

flags.DEFINE_string('log_level', 'INFO',
                    'One of VERBOSE, INFO, WARNING, ERROR. Sets logs to print '
                    'only at or above the specified level.')

flags.DEFINE_boolean('reserve_memory_for_inference_kernel', True,
                     'Normally TensorFlow preallocates the entire GPU\'s memory'
                     ' when the session is created. The inference CUDA'
                     ' executable requires some memory to run. This flag stops'
                     ' TensorFlow from allocating everything, leaving ~1GB'
                     ' for the kernel (enough for 512^3 reconstruction).'
                     ' Automatically disabled on MacOS.')


def build_model_config(dataset):
  """Creates the ModelConfig object, which contains model hyperparameters."""
  # TODO(kgenova) This needs to somehow at least support LDIF/SIF/SingleView.
  # TODO(kgenova) Add support for eval/inference.
  builder_fun_dict = {
      'ldif': hparams.build_ldif_hparams,
      'sif': hparams.build_sif_hparams,
      'sif++': hparams.build_improved_sif_hparams
  }
  model_config = experiments.ModelConfig(builder_fun_dict[FLAGS.model_type]())
  model_config.hparams.bs = FLAGS.batch_size
  model_config.train = True
  model_config.eval = False
  model_config.inference = False
  model_config.inputs['dataset'] = dataset
  model_config.inputs['split'] = FLAGS.split
  model_config.inputs['proto'] = 'ShapeNetNSSDodecaSparseLRGMediumSlimPC'
  # This function is defined by the library; we don't need it, but
  # the launcher may call it.
  model_config.wrap_optimizer = lambda x: x
  return model_config


def visualize_data(session, dataset):
  """Visualizes the dataset with two interactive visualizer windows."""
  (bounding_box_samples, depth_renders, mesh_name, near_surface_samples, grid,
   world2grid, surface_point_samples) = session.run([
       dataset.bounding_box_samples, dataset.depth_renders, dataset.mesh_name,
       dataset.near_surface_samples, dataset.grid, dataset.world2grid,
       dataset.surface_point_samples
   ])
  gaps_util.ptsview(
      [bounding_box_samples, near_surface_samples, surface_point_samples])
  mesh_name = mesh_name.decode(sys.getdefaultencoding())
  log.info(f'depth max: {np.max(depth_renders)}')
  log.info(f'Mesh name: {mesh_name}')
  assert '|' in mesh_name
  mesh_hash = mesh_name[mesh_name.find('|') + 1:]
  log.info(f'Mesh hash: {mesh_hash}')
  dyn_obj = example.InferenceExample(FLAGS.split, 'airplane', mesh_hash)

  gaps_util.gapsview(
      msh=dyn_obj.normalized_gt_mesh,
      pts=near_surface_samples[:, :3],
      grd=grid,
      world2grid=world2grid,
      grid_threshold=-0.07)


def get_model_root():
  """Finds the path to the trained model's root directory based on flags."""
  ldif_abspath = path_util.get_path_to_ldif_root()
  model_dir_is_relative = FLAGS.model_directory[0] != '/'
  if model_dir_is_relative:
    model_dir_path = os.path.join(ldif_abspath, FLAGS.model_directory)
  else:
    model_dir_path = FLAGS.model_directory
  if not os.path.isdir(model_dir_path):
    os.makedirs(model_dir_path)
  return model_dir_path


def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')
  tf.disable_v2_behavior()
  log.set_level(FLAGS.log_level)

  log.info('Making dataset...')
  if not FLAGS.dataset_directory:
    raise ValueError('A dataset directory must be provided.')
  if not os.path.isdir(FLAGS.dataset_directory):
    raise ValueError(f'No dataset directory found at {FLAGS.dataset_directory}')
  # TODO(kgenova) This batch size should match.

  dataset = local_inputs.make_dataset(