def regression(self, dpair):
model = self.model
rescale = Rescale(MODEL_CONFIG['input_size'])
try:
prev_im, curr_im = rescale({'image':
dpair['prev_im']})['image'], rescale(
{'image':
dpair['curr_im']})['image']
except:
return np.array([0, 0, 0, 0])
trans_tomodel = transforms.Compose([Normalize(), ToTensor()])
dpair = trans_tomodel({'prev_im': prev_im, 'curr_im': curr_im})
prev_im, curr_im = dpair['prev_im'], dpair['curr_im']
if self.use_gpu is True:
prev_im, curr_im = Variable(prev_im.cuda()), Variable(
curr_im.cuda())
model = model.cuda()
else:
prev_im, curr_im = Variable(prev_im), Variable(curr_im)
prev_im = prev_im[None, :, :, :]
curr_im = curr_im[None, :, :, :]
regression_bbox = model(prev_im, curr_im)
bbox = regression_bbox.data.cpu().numpy()
bbox = bbox[0, :] / MODEL_CONFIG['bbox_scale']
return bbox
def main(args):
test_ds = MnistDataset(
args.test_image_file,
args.test_label_file,
transform=transforms.Compose([ToTensor()]),
)
test_loader = torch.utils.data.DataLoader(
test_ds,
batch_size=args.batch_size,
collate_fn=collate_fn,
shuffle=False,
)
model = Net().to(device)
model.load_state_dict(torch.load(args.checkpoint))
model.eval()
predicts = []
truths = []
with torch.no_grad():
for i, sample in enumerate(test_loader):
X, Y_true = sample["X"].to(device), sample["Y"].to(device)
output = model(X)
predicts.append(torch.argmax(output, dim=1))
truths.append(Y_true)
predicts = torch.cat(predicts, dim=0)
truths = torch.cat(truths, dim=0)
acc = torch.sum(torch.eq(predicts, truths))
print("Acc: {:.4f}".format(acc / len(predicts)))
def __init__(self, mode, model=None):
"""TODO: to be defined.
:mode: TODO
:model: TODO
"""
assert isinstance(mode, str)
assert mode in pconfig.AVAILABLE_MODES
if mode in ['fully-supervised', 'val', 'test']:
assert model is None
elif mode == 'weakly-supervised':
assert model is not None
assert isinstance(model, CRAFT)
else:
raise ValueError
self._mode = mode
# NOTE [ Data Transformation ]
# 1. Must-have steps:
# 1.1. Heatmap generation (region_map, affinity_map, confidence_map)
# ----- NOTE [ Data Augmentation ] -----
# 2. A permutation of color variations:
# - Brightness
# - Contrast
# - Saturation
# - Hue
# - ...
# 3. A permutation of geometry transformations:
# - Horizontal flipping
# - Cropping
# - Scaling (input image size reduces)
# - Padding (input image size doesn't change)
# - Rotation
# - Translation
# - Shearing
# - Resizing (variable input)
# ----- END [ Data Augmentation ] -----
# 1.2. ToTensor
# 1.3. Normalization (mean-std, ...)
if self._mode in ['fully-supervised', 'weakly-supervised', 'val']:
self._generate_heatmap = GenerateHeatMap(self._mode, model)
self._resize = Resize(pconfig.TRAIN_IMAGE_SIZE)
elif self._mode == 'test':
self._resize = Resize(pconfig.TEST_IMAGE_SIZE)
else:
raise ValueError
self._to_tensor = ToTensor()
self._normalize =\
Normalize(mean=dconfig.MEAN,
std=dconfig.STD,
inplace=False)
def __init__(self, opts, tfm=ToTensor(y_long=True)):
self.rootdir = os.path.abspath(opts.rootdir)
self.slice_type = opts.slice_type
self.tfm = tfm
self.is_train = (opts.mode == 'train')
self.no_labels = opts.no_labels if hasattr(
opts, 'no_labels') else False
self.start_idx = 0
assert self.slice_type in ['SAG', 'COR', 'AX']
if self.is_train:
self.init_train()
else:
self.init_test()
def setup(self):
train_ops = []
valid_ops = []
# Image cropping for the train set
if self.hparams.random_crop:
train_ops.append(MaskedRandomCrop(self.hparams.train_inp_size))
else:
train_ops.append(CenterCrop(self.hparams.train_inp_size))
# Image cropping for the valid set
valid_ops.append(CentroidCrop(self.hparams.valid_inp_size))
# ToTensor
train_ops.append(ToTensor())
valid_ops.append(ToTensor())
# Normalize
train_ops.append(Normalize(mean=self.mean, std=self.std))
valid_ops.append(Normalize(mean=self.mean, std=self.std))
# Data split into train, valid and test sets.
indices = list(range(len(self.data)))
test_dataset = Subset(self.data, indices[:self.hparams.test_size])
valid_dataset = Subset(
self.data, indices[self.hparams.test_size:self.hparams.test_size +
self.hparams.valid_size])
train_dataset = Subset(
self.data,
indices[self.hparams.test_size + self.hparams.valid_size:])
self.datasets['train'] = RealHandsDataset(train_dataset,
Compose(train_ops))
self.datasets['valid'] = RealHandsDataset(valid_dataset,
Compose(valid_ops))
self.datasets['test'] = RealHandsDataset(test_dataset,
Compose(valid_ops))
self.datasets['mean_std'] = RealHandsDataset(self.data, ToTensor())
def preprocess(self, curr_search_reg, target_pad, curr_bb_recenter):
assert(len(curr_search_reg) == len(target_pad) == len(curr_bb_recenter)), "Length not legal..."
minibs = len(curr_search_reg)
def bboxscale(image, bbox):
h, w = image.shape[0], image.shape[1]
bbox = np.array([bbox[0]/w, bbox[1]/h, bbox[2]/w, bbox[3]/h])
bbox *= MODEL_CONFIG['bbox_scale']
return torch.from_numpy(bbox)
trans = Compose([Rescale(MODEL_CONFIG['input_size']),
Normalize(),
ToTensor()])
for i in range(minibs):
self.prev_im.append(trans(target_pad[i]))
self.curr_im.append(trans(curr_search_reg[i]))
self.curr_bb.append(bboxscale(curr_search_reg[i], curr_bb_recenter[i]))
def explore_embeddings(args):
transform = transforms.Compose([SolubilityToInt(), ToTensor()])
dataset = EmbeddingsLocalizationDataset(embeddings,
remapping,
unknown_solubility=False,
transform=transform,
descriptions_with_hash=False)
if len(
dataset[0][0].shape
) == 2: # if we have per residue embeddings they have an additional lenght dim
collate_function = padded_permuted_collate
else: # if we have reduced sequence wise embeddings use the default collate function by passing None
collate_function = None
data_loader = DataLoader(dataset,
batch_size=3,
shuffle=True,
collate_fn=collate_function)
# Needs "from models import *" to work
model = globals()[args.model_type](embeddings_dim=dataset[0][0].shape[-1],
**args.model_parameters)
with torch.no_grad():
embedding, localization, solubility, metadata = next(iter(data_loader))
# register defined hooks and run trough some example in the dataset
# model.conv1.register_forward_hook(visualize_activation_hook)
model.softmax.register_forward_hook(visualize_activation_hook)
# create mask corresponding to the zero padding used for the shorter sequecnes in the batch. All values corresponding to padding are False and the rest is True.
mask = torch.arange(metadata['length'].max())[
None, :] < metadata['length'][:, None] # [batchsize, seq_len]
model(embedding, mask)
plt.plot(torch.max(embedding[1], dim=0)[0])
plt.title('embedding max')
plt.show()
def prepare_data(self):
train_val = MnistDataset(
self.args.train_image_file,
self.args.train_label_file,
transform=transforms.Compose([ToTensor()]),
)
train_len = int(0.8 * len(train_val))
train_ds, val_ds = torch.utils.data.random_split(
train_val, [train_len, len(train_val) - train_len]
)
print("Train {}, val {}".format(len(train_ds), len(val_ds)))
self.train_loader = torch.utils.data.DataLoader(
train_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=True,
)
self.val_loader = torch.utils.data.DataLoader(
val_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=False,
)
def train():
parser = config_parser()
args = parser.parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
args.default_device = device
if args.model_type not in [
"nerf", "smpl_nerf", "append_to_nerf", "smpl", "warp",
'vertex_sphere', "smpl_estimator", "original_nerf",
'dummy_dynamic', 'image_wise_dynamic',
"append_vertex_locations_to_nerf", 'append_smpl_params'
]:
raise Exception("The model type ", args.model_type, " does not exist.")
transform = transforms.Compose([
NormalizeRGB(),
CoarseSampling(args.near, args.far, args.number_coarse_samples),
ToTensor()
])
train_dir = os.path.join(args.dataset_dir, 'train')
val_dir = os.path.join(args.dataset_dir, 'val')
if args.model_type == "nerf":
train_data = RaysFromImagesDataset(
train_dir, os.path.join(train_dir, 'transforms.json'), transform)
val_data = RaysFromImagesDataset(
val_dir, os.path.join(val_dir, 'transforms.json'), transform)
elif args.model_type == "smpl" or args.model_type == "warp":
train_data = SmplDataset(train_dir,
os.path.join(train_dir, 'transforms.json'),
args,
transform=NormalizeRGB())
val_data = SmplDataset(val_dir,
os.path.join(val_dir, 'transforms.json'),
args,
transform=NormalizeRGB())
elif args.model_type == "smpl_nerf" or args.model_type == "append_to_nerf" or args.model_type == "append_smpl_params":
train_data = SmplNerfDataset(
train_dir, os.path.join(train_dir, 'transforms.json'), transform)
val_data = SmplNerfDataset(val_dir,
os.path.join(val_dir, 'transforms.json'),
transform)
elif args.model_type == "vertex_sphere":
train_data = VertexSphereDataset(
train_dir, os.path.join(train_dir, 'transforms.json'), args)
val_data = VertexSphereDataset(
val_dir, os.path.join(val_dir, 'transforms.json'), args)
elif args.model_type == "smpl_estimator":
transform = NormalizeRGBImage()
train_data = SmplEstimatorDataset(
train_dir, os.path.join(train_dir, 'transforms.json'),
args.vertex_sphere_radius, transform)
val_data = SmplEstimatorDataset(
val_dir, os.path.join(val_dir, 'transforms.json'),
args.vertex_sphere_radius, transform)
elif args.model_type == "original_nerf":
train_data = OriginalNerfDataset(
args.dataset_dir,
os.path.join(args.dataset_dir, 'transforms_train.json'), transform)
val_data = OriginalNerfDataset(
args.dataset_dir,
os.path.join(args.dataset_dir, 'transforms_val.json'), transform)
elif args.model_type == "dummy_dynamic":
train_data = DummyDynamicDataset(
train_dir, os.path.join(train_dir, 'transforms.json'), transform)
val_data = DummyDynamicDataset(
val_dir, os.path.join(val_dir, 'transforms.json'), transform)
elif args.model_type == "append_vertex_locations_to_nerf":
train_data = DummyDynamicDataset(
train_dir, os.path.join(train_dir, 'transforms.json'), transform)
val_data = DummyDynamicDataset(
val_dir, os.path.join(val_dir, 'transforms.json'), transform)
elif args.model_type == 'image_wise_dynamic':
canonical_pose1 = torch.zeros(38).view(1, -1)
canonical_pose2 = torch.zeros(2).view(1, -1)
canonical_pose3 = torch.zeros(27).view(1, -1)
arm_angle_l = torch.tensor([np.deg2rad(10)]).float().view(1, -1)
arm_angle_r = torch.tensor([np.deg2rad(10)]).float().view(1, -1)
smpl_estimator = DummyImageWiseEstimator(canonical_pose1,
canonical_pose2,
canonical_pose3, arm_angle_l,
arm_angle_r,
torch.zeros(10).view(1, -1),
torch.zeros(69).view(1, -1))
train_data = ImageWiseDataset(
train_dir, os.path.join(train_dir, 'transforms.json'),
smpl_estimator, transform, args)
val_data = ImageWiseDataset(val_dir,
os.path.join(val_dir, 'transforms.json'),
smpl_estimator, transform, args)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batchsize,
shuffle=True,
num_workers=0)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=args.batchsize_val,
shuffle=False,
num_workers=0)
position_encoder = PositionalEncoder(args.number_frequencies_postitional,
args.use_identity_positional)
direction_encoder = PositionalEncoder(args.number_frequencies_directional,
args.use_identity_directional)
model_coarse = RenderRayNet(args.netdepth,
args.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
skips=args.skips)
model_fine = RenderRayNet(args.netdepth_fine,
args.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
skips=args.skips_fine)
if args.model_type == "smpl_nerf":
human_pose_encoder = PositionalEncoder(args.number_frequencies_pose,
args.use_identity_pose)
positions_dim = position_encoder.output_dim if args.human_pose_encoding else 1
human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1
model_warp_field = WarpFieldNet(args.netdepth_warp, args.netwidth_warp,
positions_dim * 3, human_pose_dim * 2)
solver = SmplNerfSolver(model_coarse, model_fine, model_warp_field,
position_encoder, direction_encoder,
human_pose_encoder, train_data.canonical_smpl,
args, torch.optim.Adam, torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir,
[model_coarse, model_fine, model_warp_field],
['model_coarse.pt', 'model_fine.pt', 'model_warp_field.pt'],
parser)
elif args.model_type == 'smpl':
solver = SmplSolver(model_coarse, model_fine, position_encoder,
direction_encoder, args, torch.optim.Adam,
torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w,
parser)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
elif args.model_type == 'nerf' or args.model_type == "original_nerf":
solver = NerfSolver(model_coarse, model_fine, position_encoder,
direction_encoder, args, torch.optim.Adam,
torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w,
parser)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
elif args.model_type == 'warp':
human_pose_encoder = PositionalEncoder(args.number_frequencies_pose,
args.use_identity_pose)
positions_dim = position_encoder.output_dim if args.human_pose_encoding else 1
human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1
model_warp_field = WarpFieldNet(args.netdepth_warp, args.netwidth_warp,
positions_dim * 3, human_pose_dim * 2)
human_pose_encoder = PositionalEncoder(args.number_frequencies_pose,
args.use_identity_pose)
solver = WarpSolver(model_warp_field, position_encoder,
direction_encoder, human_pose_encoder, args)
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir, [model_warp_field],
['model_warp_field.pt'], parser)
elif args.model_type == 'append_smpl_params':
human_pose_encoder = PositionalEncoder(args.number_frequencies_pose,
args.use_identity_pose)
human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1
model_coarse = RenderRayNet(
args.netdepth,
args.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args.skips,
use_directional_input=args.use_directional_input)
model_fine = RenderRayNet(
args.netdepth_fine,
args.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args.skips_fine,
use_directional_input=args.use_directional_input)
if args.load_run is not None:
model_coarse.load_state_dict(
torch.load(os.path.join(args.load_run, 'model_coarse.pt'),
map_location=torch.device(device)))
model_fine.load_state_dict(
torch.load(os.path.join(args.load_run, 'model_fine.pt'),
map_location=torch.device(device)))
print("Models loaded from ", args.load_run)
if args.siren:
model_coarse = SirenRenderRayNet(
args.netdepth,
args.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args.skips,
use_directional_input=args.use_directional_input)
model_fine = SirenRenderRayNet(
args.netdepth_fine,
args.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args.skips_fine,
use_directional_input=args.use_directional_input)
solver = AppendSmplParamsSolver(model_coarse, model_fine,
position_encoder, direction_encoder,
human_pose_encoder, args,
torch.optim.Adam, torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w,
parser)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
model_dependent = [human_pose_encoder, human_pose_dim]
inference_gif(solver.writer.log_dir, args.model_type, args, train_data,
val_data, position_encoder, direction_encoder,
model_coarse, model_fine, model_dependent)
elif args.model_type == 'append_to_nerf':
human_pose_encoder = PositionalEncoder(args.number_frequencies_pose,
args.use_identity_pose)
human_pose_dim = human_pose_encoder.output_dim if args.human_pose_encoding else 1
model_coarse = RenderRayNet(
args.netdepth,
args.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 2,
skips=args.skips,
use_directional_input=args.use_directional_input)
model_fine = RenderRayNet(
args.netdepth_fine,
args.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 2,
skips=args.skips_fine,
use_directional_input=args.use_directional_input)
solver = AppendToNerfSolver(model_coarse, model_fine, position_encoder,
direction_encoder, human_pose_encoder,
args, torch.optim.Adam, torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w,
parser)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
model_dependent = [human_pose_encoder, human_pose_dim]
inference_gif(solver.writer.log_dir, args.model_type, args, train_data,
val_data, position_encoder, direction_encoder,
model_coarse, model_fine, model_dependent)
elif args.model_type == 'append_vertex_locations_to_nerf':
model_coarse = AppendVerticesNet(args.netdepth,
args.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
6890,
additional_input_layers=1,
skips=args.skips)
model_fine = AppendVerticesNet(args.netdepth_fine,
args.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
6890,
additional_input_layers=1,
skips=args.skips_fine)
smpl_estimator = DummySmplEstimatorModel(train_data.goal_poses,
train_data.betas)
smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl"
smpl_model = smplx.create(smpl_file_name, model_type='smpl')
smpl_model.batchsize = args.batchsize
solver = AppendVerticesSolver(model_coarse, model_fine, smpl_estimator,
smpl_model, position_encoder,
direction_encoder, args,
torch.optim.Adam, torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
elif args.model_type == 'vertex_sphere':
solver = VertexSphereSolver(model_coarse, model_fine, position_encoder,
direction_encoder, args, torch.optim.Adam,
torch.nn.MSELoss())
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir, [model_coarse, model_fine],
['model_coarse.pt', 'model_fine.pt'], parser)
elif args.model_type == 'smpl_estimator':
model = SmplEstimator(human_size=len(args.human_joints))
solver = SmplEstimatorSolver(model, args, torch.optim.Adam,
torch.nn.MSELoss())
solver.train(train_loader, val_loader)
save_run(solver.writer.log_dir, [model], ['model_smpl_estimator.pt'],
parser)
elif args.model_type == "dummy_dynamic":
smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl"
smpl_model = smplx.create(smpl_file_name, model_type='smpl')
smpl_model.batchsize = args.batchsize
smpl_estimator = DummySmplEstimatorModel(train_data.goal_poses,
train_data.betas)
solver = DynamicSolver(model_fine, model_coarse, smpl_estimator,
smpl_model, position_encoder, direction_encoder,
args)
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir,
[model_coarse, model_fine, smpl_estimator],
['model_coarse.pt', 'model_fine.pt', 'smpl_estimator.pt'],
parser)
elif args.model_type == "image_wise_dynamic":
if args.load_coarse_model != None:
print("Load model..")
model_coarse.load_state_dict(
torch.load(args.load_coarse_model,
map_location=torch.device(device)))
for params in model_coarse.parameters():
params.requires_grad = False
model_coarse.eval()
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=1,
shuffle=True,
num_workers=0)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=1,
shuffle=False,
num_workers=0)
smpl_file_name = "SMPLs/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl"
smpl_model = smplx.create(smpl_file_name, model_type='smpl')
smpl_model.batchsize = args.batchsize
solver = ImageWiseSolver(model_coarse, model_fine, smpl_estimator,
smpl_model, position_encoder,
direction_encoder, args)
solver.train(train_loader, val_loader, train_data.h, train_data.w)
save_run(solver.writer.log_dir,
[model_coarse, model_fine, smpl_estimator],
['model_coarse.pt', 'model_fine.pt', 'smpl_estimator.pt'],
parser)
def setup_pipeline_dataloader(args_training, device):
position_encoder = PositionalEncoder(
args_training.number_frequencies_postitional,
args_training.use_identity_positional)
direction_encoder = PositionalEncoder(
args_training.number_frequencies_directional,
args_training.use_identity_directional)
if not args_training.inf_model_type == "append_to_nerf" and not args_training.inf_model_type == "append_smpl_params":
model_coarse = RenderRayNet(args_training.netdepth,
args_training.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
skips=args_training.skips)
model_fine = RenderRayNet(args_training.netdepth_fine,
args_training.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
skips=args_training.skips_fine)
model_coarse.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_coarse.pt"),
map_location=torch.device('cpu')))
model_coarse.eval()
model_coarse.to(device)
if os.path.exists(
os.path.join(args_training.inf_run_dir, "model_fine.pt")):
model_fine.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_fine.pt"),
map_location=torch.device('cpu')))
model_fine.eval()
model_fine.to(device)
else:
model_fine = None
transform = transforms.Compose([
NormalizeRGB(),
CoarseSampling(args_training.near, args_training.far,
args_training.number_coarse_samples),
ToTensor()
])
if args_training.inf_model_type == "smpl_nerf":
dataset = SmplNerfDataset(
args_training.inf_ground_truth_dir,
os.path.join(args_training.inf_ground_truth_dir,
'transforms.json'), transform)
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args_training.batchsize,
shuffle=False,
num_workers=0)
human_pose_encoder = PositionalEncoder(
args_training.number_frequencies_pose,
args_training.use_identity_pose)
positions_dim = position_encoder.output_dim if args_training.human_pose_encoding else 1
human_pose_dim = human_pose_encoder.output_dim if args_training.human_pose_encoding else 1
model_warp_field = WarpFieldNet(args_training.netdepth_warp,
args_training.netwidth_warp,
positions_dim * 3, human_pose_dim * 2)
model_warp_field.load_state_dict(
torch.load(
os.path.join(args_training.inf_run_dir,
"model_warp_field.pt")))
model_warp_field.eval()
pipeline = SmplNerfPipeline(model_coarse, model_fine, model_warp_field,
args_training, position_encoder,
direction_encoder, human_pose_encoder)
elif args_training.inf_model_type == "append_to_nerf":
human_pose_encoder = PositionalEncoder(
args_training.number_frequencies_pose,
args_training.use_identity_pose)
human_pose_dim = human_pose_encoder.output_dim if args_training.human_pose_encoding else 1
model_coarse = RenderRayNet(args_training.netdepth,
args_training.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 2,
skips=args_training.skips)
model_fine = RenderRayNet(args_training.netdepth_fine,
args_training.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 2,
skips=args_training.skips_fine)
model_coarse.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_coarse.pt"),
map_location=torch.device('cpu')))
model_coarse.eval()
model_fine.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_fine.pt"),
map_location=torch.device('cpu')))
model_fine.eval()
model_coarse.to(device)
model_fine.to(device)
dataset = SmplNerfDataset(
args_training.inf_ground_truth_dir,
os.path.join(args_training.inf_ground_truth_dir,
'transforms.json'), transform)
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args_training.inf_batchsize,
shuffle=False,
num_workers=0)
human_pose_encoder = PositionalEncoder(
args_training.number_frequencies_pose,
args_training.use_identity_pose)
pipeline = AppendToNerfPipeline(model_coarse, model_fine,
args_training, position_encoder,
direction_encoder, human_pose_encoder)
elif args_training.inf_model_type == "append_smpl_params":
print("Use directional input: ", args_training.use_directional_input)
human_pose_encoder = PositionalEncoder(
args_training.number_frequencies_pose,
args_training.use_identity_pose)
human_pose_dim = human_pose_encoder.output_dim if args_training.human_pose_encoding else 1
model_coarse = RenderRayNet(
args_training.netdepth,
args_training.netwidth,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args_training.skips,
use_directional_input=args_training.use_directional_input)
model_fine = RenderRayNet(
args_training.netdepth_fine,
args_training.netwidth_fine,
position_encoder.output_dim * 3,
direction_encoder.output_dim * 3,
human_pose_dim * 69,
skips=args_training.skips_fine,
use_directional_input=args_training.use_directional_input)
model_coarse.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_coarse.pt"),
map_location=torch.device('cpu')))
model_coarse.eval()
model_fine.load_state_dict(
torch.load(os.path.join(args_training.inf_run_dir,
"model_fine.pt"),
map_location=torch.device('cpu')))
model_fine.eval()
model_coarse.to(device)
model_fine.to(device)
dataset = SmplNerfDataset(
args_training.inf_ground_truth_dir,
os.path.join(args_training.inf_ground_truth_dir,
'transforms.json'), transform)
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args_training.inf_batchsize,
shuffle=False,
num_workers=0)
pipeline = AppendSmplParamsPipeline(model_coarse, model_fine,
args_training, position_encoder,
direction_encoder,
human_pose_encoder)
elif args_training.inf_model_type == "smpl":
dataset = SmplDataset(args_training.inf_ground_truth_dir,
os.path.join(args_training.inf_ground_truth_dir,
'transforms.json'),
args_training,
transform=NormalizeRGB())
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args_training.batchsize,
shuffle=False,
num_workers=0)
pipeline = SmplPipeline(model_coarse, args_training, position_encoder,
direction_encoder)
elif args_training.inf_model_type == 'nerf':
dataset = RaysFromImagesDataset(
args_training.inf_ground_truth_dir,
os.path.join(args_training.inf_ground_truth_dir,
'transforms.json'), transform)
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args_training.batchsize,
shuffle=False,
num_workers=0)
pipeline = NerfPipeline(model_coarse, model_fine, args_training,
position_encoder, direction_encoder)
return pipeline, data_loader, dataset