def __init__(self):
super().__init__()
# Define the model.
if opt.mode == 'mlp':
self.model = modules.SingleBVPNet(type=opt.model_type, final_layer_factor=1, in_features=3)
elif opt.mode == 'nerf':
self.model = modules.SingleBVPNet(type='relu', mode='nerf', final_layer_factor=1, in_features=3)
#print(datetime.datetime.now())
self.model.load_state_dict(torch.load(opt.checkpoint_path))
#print(datetime.datetime.now())
self.model.cuda()
def __init__(self,
num_instances,
latent_dim=128,
model_type='sine',
hyper_hidden_layers=1,
hyper_hidden_features=256,
hidden_num=128,
**kwargs):
super().__init__()
# We use auto-decoder framework following Park et al. 2019 (DeepSDF),
# therefore, the model consists of latent codes for each subjects and DIF-Net decoder.
# latent code embedding for training subjects
self.latent_dim = latent_dim
self.latent_codes = nn.Embedding(num_instances, self.latent_dim)
nn.init.normal_(self.latent_codes.weight, mean=0, std=0.01)
## DIF-Net
# template field
self.template_field = modules.SingleBVPNet(type=model_type,
mode='mlp',
hidden_features=hidden_num,
num_hidden_layers=3,
in_features=3,
out_features=1)
# Deform-Net
self.deform_net = modules.SingleBVPNet(type=model_type,
mode='mlp',
hidden_features=hidden_num,
num_hidden_layers=3,
in_features=3,
out_features=4)
# Hyper-Net
self.hyper_net = HyperNetwork(
hyper_in_features=self.latent_dim,
hyper_hidden_layers=hyper_hidden_layers,
hyper_hidden_features=hyper_hidden_features,
hypo_module=self.deform_net)
print(self)
def __init__(self,
in_features,
out_features,
image_resolution=None,
hypo_net_nl='sine',
encoder_nl='sine'):
super().__init__()
latent_dim = 256
self.hypo_net = modules.SingleBVPNet(out_features=out_features,
type=hypo_net_nl,
in_features=2)
def __init__(self,
in_features,
out_features,
image_resolution=None,
encoder_nl='sine'):
super().__init__()
latent_dim = 256
self.hypo_net = modules.SingleBVPNet(out_features=out_features,
type='sine',
sidelength=image_resolution,
in_features=2)
self.hyper_net = HyperNetwork(hyper_in_features=latent_dim,
hyper_hidden_layers=1,
hyper_hidden_features=256,
hypo_module=self.hypo_net)
self.set_encoder = modules.SetEncoder(in_features=in_features,
out_features=latent_dim,
num_hidden_layers=2,
hidden_features=latent_dim,
nonlinearity=encoder_nl)
print(self)
def __init__(self,
in_features,
out_features,
image_resolution=None,
partial_conv=False):
super().__init__()
latent_dim = 256
if partial_conv:
self.encoder = modules.PartialConvImgEncoder(
channel=in_features, image_resolution=image_resolution)
else:
self.encoder = modules.ConvImgEncoder(
channel=in_features, image_resolution=image_resolution)
self.hypo_net = modules.SingleBVPNet(out_features=out_features,
type='sine',
sidelength=image_resolution,
in_features=2)
self.hyper_net = HyperNetwork(hyper_in_features=latent_dim,
hyper_hidden_layers=1,
hyper_hidden_features=256,
hypo_module=self.hypo_net)
print(self)
p.add_argument('--checkpoint_path',
default=None,
help='Checkpoint to trained model.')
opt = p.parse_args()
sdf_dataset = dataio.PointCloud(opt.point_cloud_path,
on_surface_points=opt.batch_size)
dataloader = DataLoader(sdf_dataset,
shuffle=True,
batch_size=1,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu', mode='nerf', in_features=3)
else:
model = modules.SingleBVPNet(type=opt.model_type, in_features=3)
#print(datetime.datetime.now())
model.cuda()
#print(datetime.datetime.now())
# Define the loss
loss_fn = loss_functions.sdf
summary_fn = utils.write_sdf_summary
root_path = os.path.join(opt.logging_root, opt.experiment_name)
training.train(model=model,
train_dataloader=dataloader,
epochs=opt.num_epochs,
logging_root = './logs'
angle_alpha = 1.2
# Setting to plot
ckpt_path = './Deepreach_trained_checkpoints/air3D_ckpt.pth'
activation = 'sine'
times = [0.9]
time_indices_matlab = [int(time_to_plot / 0.1) + 1 for time_to_plot in times]
thetas = [1.5863] # This theta is contained in the LS computation grid.
# Initialize and load the model
model = modules.SingleBVPNet(in_features=4,
out_features=1,
type=activation,
mode='mlp',
final_layer_factor=1.,
hidden_features=512,
num_hidden_layers=3)
model.cuda()
checkpoint = torch.load(ckpt_path)
try:
model_weights = checkpoint['model']
except:
model_weights = checkpoint
model.load_state_dict(model_weights)
model.eval()
# Load the ground truth BRS data
true_BRT_path = './Deepreach_trained_checkpoints/analytical_BRT_air3D.mat'
true_data = spio.loadmat(true_BRT_path)
velocity=opt.velocity,
source_coords=source_coords)
dataloader = DataLoader(dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.mode == 'pinn':
model = modules.PINNet(out_features=2, type='tanh', mode=opt.mode)
opt.use_lbfgs = True
else:
model = modules.SingleBVPNet(out_features=2,
type=opt.model,
mode=opt.mode,
final_layer_factor=1.)
model.cuda()
# Define the loss
loss_fn = loss_functions.helmholtz_pml
summary_fn = utils.write_helmholtz_summary
root_path = os.path.join(opt.logging_root, opt.experiment_name)
training.train(model=model,
train_dataloader=dataloader,
epochs=opt.num_epochs,
lr=opt.lr,
steps_til_summary=opt.steps_til_summary,
img_dataset = dataio.BSD500ImageDataset(in_folder='../data/BSD500/train',
idx_to_sample=[19])
coord_dataset = dataio.Implicit2DWrapper(img_dataset,
sidelength=256,
compute_diff='gradients')
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.model_type == 'sine' or opt.model_type == 'relu' or opt.model_type == 'tanh' or opt.model_type == 'softplus':
model = modules.SingleBVPNet(type=opt.model_type,
mode='mlp',
sidelength=(256, 256))
elif opt.model_type == 'rbf' or opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu',
mode=opt.model_type,
sidelength=(256, 256))
else:
raise NotImplementedError
model.cuda()
# Define the loss & summary functions
loss_fn = loss_functions.gradients_mse
summary_fn = utils.write_gradients_summary
root_path = os.path.join(opt.logging_root, opt.experiment_name)
if opt.counter_start == -1:
opt.counter_start = opt.checkpoint_toload
if opt.counter_end == -1:
opt.counter_end = opt.num_epochs
dataset = dataio.ReachabilityAir3DSource(numpoints=65000, collisionR=opt.collisionR, velocity=opt.velocity,
omega_max=opt.omega_max, pretrain=opt.pretrain, tMin=opt.tMin,
tMax=opt.tMax, counter_start=opt.counter_start, counter_end=opt.counter_end,
pretrain_iters=opt.pretrain_iters, seed=opt.seed,
angle_alpha=opt.angle_alpha,
num_src_samples=opt.num_src_samples)
dataloader = DataLoader(dataset, shuffle=True, batch_size=opt.batch_size, pin_memory=True, num_workers=0)
model = modules.SingleBVPNet(in_features=4, out_features=1, type=opt.model, mode=opt.mode,
final_layer_factor=1., hidden_features=opt.num_nl, num_hidden_layers=opt.num_hl)
model.cuda()
# Define the loss
loss_fn = loss_functions.initialize_hji_air3D(dataset, opt.minWith)
root_path = os.path.join(opt.logging_root, opt.experiment_name)
def val_fn(model, ckpt_dir, epoch):
# Time values at which the function needs to be plotted
times = [0., 0.5*(opt.tMax - 0.1), (opt.tMax - 0.1)]
num_times = len(times)
# Theta slices to be plotted
thetas = [-math.pi, -0.5*math.pi, 0., 0.5*math.pi, math.pi]
num_thetas = len(thetas)
coord_dataset = dataio.Implicit2DWrapper(img_dataset,
sidelength=512,
compute_diff='all')
image_resolution = (512, 512)
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.model_type == 'sine' or opt.model_type == 'relu' or opt.model_type == 'tanh' or opt.model_type == 'selu' or opt.model_type == 'elu'\
or opt.model_type == 'softplus':
model = modules.SingleBVPNet(type=opt.model_type,
mode='mlp',
sidelength=image_resolution)
elif opt.model_type == 'rbf' or opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu',
mode=opt.model_type,
sidelength=image_resolution)
else:
raise NotImplementedError
model.cuda()
root_path = os.path.join(opt.logging_root, opt.experiment_name)
# Define the loss
loss_fn = partial(loss_functions.image_mse, None)
summary_fn = partial(utils.write_image_summary, image_resolution)
p.add_argument('--model_type', type=str, default='sine',
help='Options currently are "sine" (all sine activations), "relu" (all relu activations,'
'"nerf" (relu activations and positional encoding as in NeRF), "rbf" (input rbf layer, rest relu),'
'and in the future: "mixed" (first layer sine, other layers tanh)')
p.add_argument('--checkpoint_path', required=True, help='Checkpoint to trained model.')
opt = p.parse_args()
audio_dataset = dataio.AudioFile(filename=opt.gt_wav_path)
coord_dataset = dataio.ImplicitAudioWrapper(audio_dataset)
dataloader = DataLoader(coord_dataset, shuffle=True, batch_size=1, pin_memory=True, num_workers=0)
# Define the model and load in checkpoint path
if opt.model_type == 'sine' or opt.model_type == 'relu' or opt.model_type == 'tanh':
model = modules.SingleBVPNet(type=opt.model_type, mode='mlp', in_features=1)
elif opt.model_type == 'rbf' or opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu', mode=opt.model_type, fn_samples=len(audio_dataset.data), in_features=1)
else:
raise NotImplementedError
model.load_state_dict(torch.load(opt.checkpoint_path))
model.cuda()
root_path = os.path.join(opt.logging_root, opt.experiment_name)
utils.cond_mkdir(root_path)
# Get ground truth and input data
model_input, gt = next(iter(dataloader))
model_input = {key: value.cuda() for key, value in model_input.items()}
gt = {key: value.cuda() for key, value in gt.items()}
img_dataset,
sidelength=(img_dataset[0].size[1], img_dataset[0].size[0]),
compute_diff='all')
image_resolution = (img_dataset[0].size[1], img_dataset[0].size[0])
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.model_type == 'sine' or opt.model_type == 'relu' or opt.model_type == 'tanh':
model = modules.SingleBVPNet(type=opt.model_type,
mode='mlp',
out_features=img_dataset.img_channels,
sidelength=image_resolution,
downsample=opt.downsample)
elif opt.model_type == 'rbf' or opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu',
mode=opt.model_type,
out_features=img_dataset.img_channels,
sidelength=image_resolution,
downsample=opt.downsample)
else:
raise NotImplementedError
model.cuda()
root_path = os.path.join(opt.logging_root, opt.experiment_name)
if opt.mask_path:
numEvaders=opt.numEvaders,
pretrain_iters=opt.pretrain_iters,
angle_alpha=opt.angle_alpha,
time_alpha=opt.time_alpha,
num_src_samples=opt.num_src_samples)
dataloader = DataLoader(dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
model = modules.SingleBVPNet(in_features=(opt.numEvaders + 1) * 3 + 1,
out_features=1,
type=opt.model,
mode=opt.mode,
final_layer_factor=1.,
hidden_features=512,
num_hidden_layers=opt.num_hl)
model.cuda()
# Define the loss
loss_fn = loss_functions.initialize_hji_MultiVehicleCollisionNE(
dataset, opt.minWith)
root_path = os.path.join(opt.logging_root, opt.experiment_name)
def val_fn(model, ckpt_dir, epoch):
return
vid_dataset = dataio.Video(video_path)
coord_dataset = dataio.Implicit3DWrapper(vid_dataset,
sidelength=vid_dataset.shape,
sample_fraction=opt.sample_frac)
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if opt.model_type == 'sine' or opt.model_type == 'relu' or opt.model_type == 'tanh':
model = modules.SingleBVPNet(type=opt.model_type,
in_features=3,
out_features=vid_dataset.channels,
mode='mlp',
hidden_features=1024,
num_hidden_layers=3)
elif opt.model_type == 'rbf' or opt.model_type == 'nerf':
model = modules.SingleBVPNet(type='relu',
in_features=3,
out_features=vid_dataset.channels,
mode=opt.model_type)
else:
raise NotImplementedError
model.cuda()
root_path = os.path.join(opt.logging_root, opt.experiment_name)
# Define the loss
loss_fn = partial(loss_functions.image_mse, None)
img_filepath1 = '/media/data4e/jnmartel/neural_bvps/gizeh.jpg'
img_filepath2 = '/media/data4e/jnmartel/neural_bvps/bear.jpg'
is_color = False
coord_dataset = dataio.CompositeGradients(img_filepath1,
img_filepath2,
is_color=is_color,
sidelength=512)
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
if not is_color:
model = modules.SingleBVPNet(type=opt.model_type, final_layer_factor=1)
loss_fn = loss_functions.gradients_mse
else:
model = modules.SingleBVPNet(out_features=3,
type=opt.model_type,
final_layer_factor=1)
loss_fn = loss_functions.gradients_color_mse
model.cuda()
summary_fn = utils.write_gradcomp_summary
root_path = os.path.join(opt.logging_root, opt.experiment_name)
training.train(model=model,
train_dataloader=dataloader,
epochs=opt.num_epochs,
poss['2E'] = [(0.43, 0.33)]
thetas['2E'] = [-2.44]
# Theta of the ego vehicle
ego_vehicle_theta = [-2.54]
# Time at which the sets should be plotted
time = 1.0
# Number of slices to plot
num_slices = len(poss['1E'])
# Load the model
model = modules.SingleBVPNet(in_features=10,
out_features=1,
type='sine',
mode='mlp',
final_layer_factor=1.,
hidden_features=512,
num_hidden_layers=3)
model.cuda()
checkpoint = torch.load(ckpt_path)
try:
model_weights = checkpoint['model']
except:
model_weights = checkpoint
model.load_state_dict(model_weights)
model.eval()
# Load the 2 vehicle model
model_1P = modules.SingleBVPNet(in_features=7,
out_features=1,
dataset = dataio.InverseHelmholtz(source_coords,
rec_coords,
rec_val,
sidelength=115,
pretrain=opt.pretrain)
dataloader = DataLoader(dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
# Define the model.
N_src = source_coords.shape[0]
model = modules.SingleBVPNet(in_features=2,
out_features=2 * N_src + 1,
type=opt.model,
final_layer_factor=1.)
if opt.load_model is not None:
model.load_state_dict(torch.load(opt.load_model))
model.cuda()
# Define the loss
loss_fn = loss_functions.helmholtz_pml
summary_fn = utils.write_helmholtz_summary
root_path = os.path.join(opt.logging_root, opt.experiment_name)
training.train(model=model,
train_dataloader=dataloader,
