def define_model(self):
opts = self.opts
# ----------
# Options
# ----------
self.symmetric = opts.symmetric
anno_sfm_path = osp.join(opts.cub_cache_dir, 'sfm', 'anno_train.mat')
anno_sfm = sio.loadmat(anno_sfm_path,
struct_as_record=False,
squeeze_me=True)
if opts.sphere_initial:
sfm_mean_shape = mesh.create_sphere(3)
else:
sfm_mean_shape = (np.transpose(anno_sfm['S']),
anno_sfm['conv_tri'] - 1)
img_size = (opts.img_size, opts.img_size)
self.model = mesh_net.MeshNet(img_size,
opts,
nz_feat=opts.nz_feat,
num_kps=opts.num_kps,
sfm_mean_shape=sfm_mean_shape)
if opts.num_pretrain_epochs > 0:
self.load_network(self.model, 'pred', opts.num_pretrain_epochs)
self.model = self.model.cuda(device=opts.gpu_id)
# Data structures to use for triangle priors.
edges2verts = self.model.edges2verts
# B x E x 4
edges2verts = np.tile(np.expand_dims(edges2verts, 0),
(opts.batch_size, 1, 1))
self.edges2verts = Variable(
torch.LongTensor(edges2verts).cuda(device=opts.gpu_id),
requires_grad=False)
# For renderering.
faces = self.model.faces.view(1, -1, 3)
self.faces = faces.repeat(opts.batch_size, 1, 1)
# opts.renderer = "smr"
self.renderer = NeuralRenderer(
opts.img_size) if opts.renderer == "nmr" else SoftRenderer(
opts.img_size)
self.renderer_predcam = NeuralRenderer(
opts.img_size) if opts.renderer == "nmr" else SoftRenderer(
opts.img_size) #for camera loss via projection
# Need separate NMR for each fwd/bwd call.
if opts.texture:
self.tex_renderer = NeuralRenderer(
opts.img_size) if opts.renderer == "nmr" else SoftRenderer(
opts.img_size)
# Only use ambient light for tex renderer
self.tex_renderer.ambient_light_only()
# For visualization
self.vis_rend = bird_vis.VisRenderer(opts.img_size,
faces.data.cpu().numpy())
return
def __init__(self, opts):
self.opts = opts
self.symmetric = opts.symmetric
img_size = (opts.img_size, opts.img_size)
print('Setting up model..')
self.model = mesh_net.MeshNet(img_size, opts, nz_feat=opts.nz_feat)
self.load_network(self.model, 'pred', self.opts.num_train_epoch)
self.model.eval()
self.model = self.model.cuda(device=self.opts.gpu_id)
self.renderer = NeuralRenderer(opts.img_size)
if opts.texture:
self.tex_renderer = NeuralRenderer(opts.img_size)
# Only use ambient light for tex renderer
self.tex_renderer.ambient_light_only()
if opts.use_sfm_ms:
anno_sfm_path = osp.join(opts.cub_cache_dir, 'sfm',
'anno_testval.mat')
anno_sfm = sio.loadmat(anno_sfm_path,
struct_as_record=False,
squeeze_me=True)
if opts.sphere_initial:
sfm_mean_shape, sfm_face = mesh.create_sphere(3)
else:
sfm_mean_shape = torch.Tensor(np.transpose(
anno_sfm['S'])).cuda(device=opts.gpu_id)
self.sfm_mean_shape = Variable(sfm_mean_shape,
requires_grad=False)
self.sfm_mean_shape = self.sfm_mean_shape.unsqueeze(0).repeat(
opts.batch_size, 1, 1)
sfm_face = torch.LongTensor(anno_sfm['conv_tri'] -
1).cuda(device=opts.gpu_id)
self.sfm_face = Variable(sfm_face, requires_grad=False)
faces = self.sfm_face.view(1, -1, 3)
else:
# For visualization
faces = self.model.faces.view(1, -1, 3)
self.faces = faces.repeat(opts.batch_size, 1, 1)
self.vis_rend = bird_vis.VisRenderer(opts.img_size,
faces.data.cpu().numpy())
self.vis_rend.set_bgcolor([1., 1., 1.])
self.resnet_transform = torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
from utils.viz_flow import viz_flow
# H x W x 2
flow = convert2np(flow_img)
# viz_flow expects the top left to be zero.
# Conver to image coord
flow = (flow + 1) * 0.5 * img_size
flow_img = viz_flow(flow[:, :, 1], flow[:, :, 0])
return flow_img
if __name__ == '__main__':
# Test vis_vert2kp:
from utils import mesh
verts, faces = mesh.create_sphere()
num_kps = 15
num_vs = verts.shape[0]
ind = np.random.randint(0, num_vs, num_vs)
dists = np.stack([
np.linalg.norm(verts - verts[np.random.randint(0, num_vs)], axis=1)
for k in range(num_kps)
])
vert2kp = np.exp(-.5 * (dists) / (np.random.rand(num_kps, 1) + 0.4))
vert2kp = vert2kp / vert2kp.sum(1).reshape(-1, 1)
vis_vert2kp(verts, vert2kp, faces)
def __init__(self,
input_shape,
opts,
nz_feat=100,
num_kps=15,
sfm_mean_shape=None):
# Input shape is H x W of the image.
super(MeshNet, self).__init__()
self.opts = opts
self.pred_texture = opts.texture
self.symmetric = opts.symmetric
self.symmetric_texture = opts.symmetric_texture
# Mean shape.
verts, faces = mesh.create_sphere(opts.subdivide)
num_verts = verts.shape[0]
if self.symmetric:
verts, faces, num_indept, num_sym, num_indept_faces, num_sym_faces = mesh.make_symmetric(
verts, faces)
if sfm_mean_shape is not None:
verts = geom_utils.project_verts_on_mesh(
verts, sfm_mean_shape[0], sfm_mean_shape[1])
num_sym_output = num_indept + num_sym
if opts.only_mean_sym:
print('Only the mean shape is symmetric!')
self.num_output = num_verts
else:
self.num_output = num_sym_output
self.num_sym = num_sym
self.num_indept = num_indept
self.num_indept_faces = num_indept_faces
self.num_sym_faces = num_sym_faces
# mean shape is only half.
self.mean_v = nn.Parameter(torch.Tensor(verts[:num_sym_output]))
# Needed for symmetrizing..
self.flip = Variable(torch.ones(1, 3).cuda(), requires_grad=False)
self.flip[0, 0] = -1
else:
if sfm_mean_shape is not None:
verts = geom_utils.project_verts_on_mesh(
verts, sfm_mean_shape[0], sfm_mean_shape[1])
self.mean_v = nn.Parameter(torch.Tensor(verts))
self.num_output = num_verts
verts_np = verts
faces_np = faces
self.faces = Variable(torch.LongTensor(faces).cuda(),
requires_grad=False)
self.edges2verts = mesh.compute_edges2verts(verts, faces)
vert2kp_init = torch.Tensor(
np.ones((num_kps, num_verts)) / float(num_verts))
# Remember initial vert2kp (after softmax)
self.vert2kp_init = torch.nn.functional.softmax(Variable(
vert2kp_init.cuda(), requires_grad=False),
dim=1)
self.vert2kp = nn.Parameter(vert2kp_init)
self.encoder = Encoder(input_shape, n_blocks=4, nz_feat=nz_feat)
self.code_predictor = CodePredictor(nz_feat=nz_feat,
num_verts=self.num_output)
if self.pred_texture:
if self.symmetric_texture:
num_faces = self.num_indept_faces + self.num_sym_faces
else:
num_faces = faces.shape[0]
uv_sampler = mesh.compute_uvsampler(verts_np,
faces_np[:num_faces],
tex_size=opts.tex_size)
# F' x T x T x 2
uv_sampler = Variable(torch.FloatTensor(uv_sampler).cuda(),
requires_grad=False)
# B x F' x T x T x 2
uv_sampler = uv_sampler.unsqueeze(0).repeat(
self.opts.batch_size, 1, 1, 1, 1)
img_H = int(2**np.floor(np.log2(
np.sqrt(num_faces) * opts.tex_size)))
img_W = 2 * img_H
self.texture_predictor = TexturePredictorUV(
nz_feat,
uv_sampler,
opts,
img_H=img_H,
img_W=img_W,
predict_flow=True,
symmetric=opts.symmetric_texture,
num_sym_faces=self.num_sym_faces)
nb.net_init(self.texture_predictor)
# NOTE: jz, uv_sampler for DIB-R
self.uv_sampler = uv_sampler