def __init__(self, filename_obj, cfg):
super(Reconstructor, self).__init__()
feature_dim = cfg['Reconstruction']['f_dim']
img_size = cfg['input_size']
if cfg['Reconstruction']['encoder_type'] == 'resnet':
self.encoder = ResnetEncoder(c_dim=feature_dim)
elif ['Reconstruction']['encoder_type'] == 'simple':
self.encoder = SimpleEncoder(dim_out=feature_dim, im_size=img_size)
else:
raise Exception('encoder type must be resnet or simple')
self.color_rec = cfg['Reconstruction']['color_rec']
self.decoder = Decoder(filename_obj,
self.color_rec,
color_num=cfg['Reconstruction']['color_num'])
if self.color_rec:
self.renderer = sr.SoftRenderer(
image_size=img_size,
sigma_val=cfg['SoftRender']['SIGMA_VAL'],
aggr_func_rgb='softmax',
camera_mode='look_at',
viewing_angle=15,
dist_eps=1e-10)
else:
self.renderer = sr.SoftRenderer(
image_size=img_size,
sigma_val=cfg['SoftRender']['SIGMA_VAL'],
aggr_func_rgb='hard',
camera_mode='look_at',
viewing_angle=15,
dist_eps=1e-10)
self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base,
self.decoder.faces)
self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def define_criterion(self):
# shape objectives
self.mask_loss_fn = loss_utils.neg_iou_loss
if(opts.multi_gpu):
verts = self.model.module.get_mean_shape().cpu()
faces = self.model.module.faces.cpu()
else:
verts = self.model.get_mean_shape().cpu()
faces = self.model.faces.cpu()
self.laplacian_loss_fn = sr.LaplacianLoss(verts, faces).cuda()
self.flatten_loss_fn = sr.FlattenLoss(faces).cuda()
if(opts.multi_gpu):
self.laplacian_loss_fn = torch.nn.DataParallel(self.laplacian_loss_fn)
self.flatten_loss_fn = torch.nn.DataParallel(self.flatten_loss_fn)
# shape constraints
self.deform_reg_fn = loss_utils.deform_l2reg
self.ori_reg_fn = loss_utils.sym_reg
self.gan_loss_fn = torch.nn.functional.binary_cross_entropy_with_logits
# texture objectives
if self.opts.use_texture:
self.texture_loss = loss_utils.PerceptualTextureLoss()
self.texture_dt_loss_fn = loss_utils.texture_dt_loss
self.texture_cycle_fn = loss_utils.TexCycle(int(opts.batch_size/opts.gpu_num))
self.texture_cycle_fn = self.texture_cycle_fn.cuda()
if(opts.multi_gpu):
self.texture_cycle_fn = torch.nn.DataParallel(self.texture_cycle_fn)
def __init__(self, template_path, bg_path):
super(Model, self).__init__()
# set template mesh
self.template_mesh = sr.Mesh.from_obj(template_path)
self.bg_mesh = sr.Mesh.from_obj(bg_path)
self.register_buffer('vertices', self.template_mesh.vertices * 0.5)
self.register_buffer('faces', self.template_mesh.faces)
self.register_buffer(
'bg_vertices',
self.bg_mesh.vertices * 3 + torch.tensor([0, -0.5, 0]).repeat(
1, self.bg_mesh.vertices.shape[1], 1).cuda())
self.register_buffer('bg_faces', self.bg_mesh.faces)
# Make the object black and background white
textures = torch.cat([
torch.tensor([0., 0., 0.]).repeat(
1, self.template_mesh.vertices.shape[1], 1),
torch.tensor([1., 1., 1.]).repeat(
1, self.bg_mesh.vertices.shape[1], 1),
],
dim=1)
self.register_buffer('all_textures', textures)
# optimize for displacement map and center
self.register_parameter(
'displace',
nn.Parameter(torch.zeros_like(self.template_mesh.vertices)))
self.register_parameter('center', nn.Parameter(torch.zeros(1, 1, 3)))
# define Laplacian and flatten geometry constraints
self.laplacian_loss = sr.LaplacianLoss(self.vertices[0].cpu(),
self.faces[0].cpu())
self.flatten_loss = sr.FlattenLoss(self.faces[0].cpu())
def __init__(self, template_path):
super(Decoder, self).__init__()
self.template_mesh = sr.Mesh.from_obj(template_path)
self.register_buffer('vertices', self.template_mesh.vertices * 0.5)
self.register_buffer('faces', self.template_mesh.faces)
self.laplacian_loss = sr.LaplacianLoss(self.vertices[0].cpu(), self.faces[0].cpu())
self.flatten_loss = sr.FlattenLoss(self.faces[0].cpu())
self.fc1 = nn.Sequential(
nn.Linear(512, 2*2*512),
nn.ReLU()
)
self.deconv1 = nn.Sequential(
nn.ConvTranspose2d(512, 256, 3, stride=2),
nn.ReLU()
)
self.deconv2 = nn.Sequential(
nn.ConvTranspose2d(256, 128, 3, stride=2),
nn.ReLU()
)
self.deconv3 = nn.Sequential(
nn.ConvTranspose2d(128, 64, 3, stride=2),
nn.ReLU()
)
self.fc2 = nn.Sequential(
nn.Linear(23*23*64, 2048),
nn.ReLU())
self.fc_center = nn.Linear(2048, 3)
self.fc_displace = nn.Linear(2048, self.template_mesh.num_vertices * 3)
def __init__(self, filename_obj, args):
super(Model, self).__init__()
self.encoder = VPLEncoder()
self.decoder = VPLDecoder(filename_obj)
self.renderer = sr.SoftRenderer(image_size=args.image_size, sigma_val=args.sigma_val,
aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15,
dist_eps=1e-10)
self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base, self.decoder.faces)
self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def __init__(self, filename_obj, cfg):
super(Unconditional_Generator, self).__init__()
z_dim = cfg['NormalGan']['G']['z_dim']
self.decoder = Decoder(filename_obj, color_rec=False, dim_in=z_dim)
self.renderer = sr.SoftRenderer(
image_size=cfg['input_size'],
sigma_val=cfg['SoftRender']['SIGMA_VAL'],
aggr_func_rgb='hard',
camera_mode='look_at',
viewing_angle=15,
dist_eps=1e-10)
self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base,
self.decoder.faces)
self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def define_criterion(self):
self.projection_loss = loss_utils.kp_l2_loss
self.mask_loss_fn = torch.nn.MSELoss()
self.entropy_loss = loss_utils.entropy_loss
self.deform_reg_fn = loss_utils.deform_l2reg
self.camera_loss = loss_utils.camera_loss
self.triangle_loss_fn = loss_utils.LaplacianLoss(self.faces)
if self.opts.texture:
self.texture_loss = loss_utils.PerceptualTextureLoss()
self.texture_dt_loss_fn = loss_utils.texture_dt_loss
if self.opts.add_smr_loss:
self.flatten_loss_fn = sr.FlattenLoss(self.model.faces)
self.ori_reg_fn = loss_utils.sym_reg
def __init__(self, template_path):
super(Decoder, self).__init__()
self.template_mesh = sr.Mesh.from_obj(template_path)
self.register_buffer('vertices', self.template_mesh.vertices * 0.5)
self.register_buffer('faces', self.template_mesh.faces)
self.laplacian_loss = sr.LaplacianLoss(self.vertices[0].cpu(),
self.faces[0].cpu())
self.flatten_loss = sr.FlattenLoss(self.faces[0].cpu())
self.fc1 = nn.Sequential(nn.Linear(512, 1024), nn.ReLU())
self.fc2 = nn.Sequential(nn.Linear(1024, 2048), nn.ReLU())
self.fc_center = nn.Linear(2048, 3)
self.fc_displace = nn.Linear(2048, self.template_mesh.num_vertices * 3)
def __init__(self, args):
super(VSLModel, self).__init__()
self.batch_size = args.batch_size
self.inf_model = InfModel(args=args)
self.gen_model = GenModel(args=args)
self.image_decoder = ImageDecoder(args=args)
self.renderer = sr.SoftRenderer(image_size=args.image_size,
sigma_val=args.sigma_val,
aggr_func_rgb='hard',
camera_mode='look_at',
viewing_angle=15,
dist_eps=1e-10)
self.laplacian_loss = sr.LaplacianLoss(self.gen_model.vertices_base,
self.gen_model.faces)
self.flatten_loss = sr.FlattenLoss(self.gen_model.faces)
def __init__(self, filename_obj, args):
super(Model, self).__init__()
# auto-encoder
self.encoder = Encoder(im_size=args.image_size)
self.decoder = Decoder(filename_obj)
# renderer
self.transform = sr.LookAt(viewing_angle=15)
self.lighting = sr.Lighting()
self.rasterizer = sr.SoftRasterizer(image_size=args.image_size, sigma_val=args.sigma_val,
aggr_func_rgb='hard', aggr_func_alpha='prod', dist_eps=1e-10)
# mesh regularizer
self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base, self.decoder.faces)
self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def define_criterion(self):
# shape objectives
self.mask_loss_fn = loss_utils.MultiMaskLoss(opts.image_size,
opts.renderer_type,
opts.num_hypo_cams)
self.mask_loss_fn = torch.nn.DataParallel(self.mask_loss_fn)
verts = self.mean_shape.cpu()
faces = self.faces[0].cpu()
self.laplacian_loss_fn = sr.LaplacianLoss(verts, faces).cuda()
self.flatten_loss_fn = sr.FlattenLoss(faces).cuda()
if (opts.multi_gpu):
self.laplacian_loss_fn = torch.nn.DataParallel(
self.laplacian_loss_fn)
self.flatten_loss_fn = torch.nn.DataParallel(self.flatten_loss_fn)
self.gan_loss_fn = torch.nn.functional.binary_cross_entropy_with_logits
self.deform_reg_fn = loss_utils.deform_l2reg
# texture objectives
self.texture_loss_fn = loss_utils.MultiTextureLoss(
int(opts.batch_size / opts.gpu_num))
self.texture_loss_fn = torch.nn.DataParallel(self.texture_loss_fn)
# semantic correspondence constraints
self.corr_loss_fn = loss_utils.CorrLossChamfer(opts.stemp_path,
opts.image_size)
if (opts.multi_gpu):
self.part_loss_fn = loss_utils.part_matching_loss(
opts.stemp_path,
self.model.module.uv_sampler[0].unsqueeze(0),
self.model.module.texture_predictor.num_sym_faces,
batch_size=int(opts.batch_size / opts.gpu_num))
else:
self.part_loss_fn = loss_utils.part_matching_loss(
opts.stemp_path,
self.model.uv_sampler[0].unsqueeze(0),
self.model.texture_predictor.num_sym_faces,
batch_size=int(opts.batch_size / opts.gpu_num))
self.part_loss_fn = self.part_loss_fn.cuda()
self.part_loss_fn = torch.nn.DataParallel(self.part_loss_fn)
def __init__(self, template_path):
super(Model, self).__init__()
# set template mesh
self.template_mesh = sr.Mesh.from_obj(template_path)
self.register_buffer('vertices', self.template_mesh.vertices)
self.register_buffer('faces', self.template_mesh.faces)
self.register_buffer('textures', self.template_mesh.textures)
# optimize for displacement map and center
self.register_parameter(
'displace',
nn.Parameter(torch.zeros_like(self.template_mesh.vertices)))
self.register_parameter('center', nn.Parameter(torch.zeros(1, 1, 3)))
# define Laplacian and flatten geometry constraints
self.laplacian_loss = sr.LaplacianLoss(self.vertices[0].cpu(),
self.faces[0].cpu())
self.flatten_loss = sr.FlattenLoss(self.faces[0].cpu())
def __init__(self, template_path):
super(Decoder, self).__init__()
self.template_mesh = sr.Mesh.from_obj(template_path)
self.register_buffer('vertices', self.template_mesh.vertices * 0.55)
self.register_buffer('faces', self.template_mesh.faces)
self.laplacian_loss = sr.LaplacianLoss(self.vertices[0].cpu(),
self.faces[0].cpu())
self.flatten_loss = sr.FlattenLoss(self.faces[0].cpu())
self.fc1 = nn.Sequential(nn.Linear(15 * 15 * 256, 5000), nn.ReLU())
# self.fc2 = nn.Sequential(
# nn.Linear(4096, 4096),
# nn.ReLU()
# )
# self.fc3 = nn.Sequential(
# nn.Linear(4096, 4096),
# nn.ReLU()
# )
self.fc_center = nn.Linear(5000, 3)
self.fc_displace = nn.Linear(5000, self.template_mesh.num_vertices * 3)
