viewer_pos, viewer_look = du.get_random_viewer_info(self.c_sl)
dir_vec = du.random_dir_vec_tensor()
target = get_shape_dir_target(viewer_pos, dir_vec, c_sizes, s_sizes,
self.c_sl, self.max_shift)
context = su.get_dense_array_from_sl(context_sparse, self.c_sl,
self.useid)
seg = su.get_dense_array_from_sl(seg_sparse, self.s_sl, self.useid)
return {
"context": torch.from_numpy(context),
"seg": torch.from_numpy(seg),
"target": target,
"viewer_pos": viewer_pos,
"dir_vec": dir_vec,
}
def __getitem__(self, index):
return self._get_example()
def __len__(self):
return self.num_examples
if __name__ == "__main__":
from visualization_utils import GeoscorerDatasetVisualizer
dataset = SegmentContextShapeDirData(nexamples=3)
vis = GeoscorerDatasetVisualizer(dataset)
for n in range(len(dataset)):
vis.visualize()
dataset, dataloader = tu.setup_dataset_and_loader(opts)
tms = tu.get_context_segment_trainer_modules(opts,
opts["checkpoint"],
backup=opts["backup"],
verbose=True)
# The context and seg net were already moved
if opts["cuda"] == 1:
tms["score_module"].cuda()
tms["lfn"].cuda()
# Setup visualization
vis = None
if opts["visualize_epochs"]:
from visualization_utils import GeoscorerDatasetVisualizer
vis = GeoscorerDatasetVisualizer(dataset)
vis.set_model(tms, opts)
run_visualization(vis,
tms,
opts,
checkpoint_path=None,
num_examples=2,
tsleep=1,
loadpath=False)
# Run training
for i in range(opts["nepoch"]):
train_loss, train_error = train_epoch(tms, dataloader, opts)
tu.pretty_log("train loss {:<5.4f} error {:<5.2f} {}".format(
train_loss, train_error * 100, i))
if opts["checkpoint"] != "":
if opts.dataset_type == "pair":
dataset = ShapePairData(
nexamples=opts.n_examples,
ground_type=opts.ground_type,
shape_type=opts.shape_type,
max_shift=opts.max_shift,
fixed_size=opts.fixed_size,
random_ground_height=opts.random_ground_height,
)
else:
dataset = ShapePieceData(
nexamples=opts.n_examples,
ground_type=opts.ground_type,
random_ground_height=opts.random_ground_height,
)
vis = GeoscorerDatasetVisualizer(dataset)
use_model = opts.checkpoint is not None
if use_model:
import training_utils as tu
tms = tu.get_context_segment_trainer_modules(
vars(opts), opts.checkpoint, backup=False, verbose=True, use_new_opts=False
)
if tms["opts"]["cuda"] == 1:
tms["score_module"].cuda()
tms["lfn"].cuda()
vis.set_model(tms, opts=tms["opts"])
for n in range(len(dataset)):
vis.visualize(use_model=use_model, verbose=True)