atlasnet.py

import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo
import data
import util
import chamfer

# network
class AtlasNet(torch.nn.Module):

	def __init__(self,opt,eval_enc=False,eval_dec=False):
		super(AtlasNet,self).__init__()
		# define UV
		self.UV_sphere,self.faces_sphere = data.get_icosahedron(opt)
		self.UV_regular,self.faces_regular = self.get_regular_patch_grid(opt)
		self.faces_regular = self.duplicate_faces_original(opt)
		# define and load pretrained weights
		self.define_weights(opt)
		if opt.pretrained_dec is not None:
			self.load_pretrained_decoder(opt)
		for p in self.encoder.parameters(): p.requires_grad_(not eval_enc)
		for p in self.decoder.parameters(): p.requires_grad_(not eval_dec)
		(self.encoder.eval if eval_enc else self.encoder.train)()
		(self.decoder.eval if eval_dec else self.decoder.train)()

	def define_weights(self,opt):
		embed_size = 1024
		self.encoder = resnet18(pretrained=opt.imagenet_enc,num_classes=embed_size)
		code_size = embed_size+3 if opt.sphere else embed_size+2
		self.decoder = torch.nn.ModuleList([PointGenCon(code_size=code_size) for _ in range(opt.num_prim)])
		self.encoder = self.encoder.to(opt.device)
		self.decoder = self.decoder.to(opt.device)

	def load_pretrained_decoder(self,opt):
		print(util.magenta("loading pretrained decoder ({})...".format(opt.pretrained_dec)))
		weight_dict = torch.load(opt.pretrained_dec,map_location=opt.device)
		# remove "decoder/" prefix in dictionary
		decoder_weight_dict = {k[8:]: weight_dict[k] for k in weight_dict if "decoder" in k}
		self.decoder.load_state_dict(decoder_weight_dict)

	def decoder_forward(self,opt,code,regular=False):
		batch_size = code.shape[0]
		points_list = []
		for p in range(opt.num_prim):
			if opt.sphere:
				UV = self.UV_sphere.repeat(batch_size,1,1).permute(0,2,1)
			else:
				if regular:
					UV = self.UV_regular.repeat(batch_size,1,1).permute(0,2,1)
				else:
					UV = torch.rand(batch_size,2,opt.num_points,device=opt.device)
			concat = torch.cat([UV,code[...,None].repeat(1,1,UV.shape[2])],dim=1)
			points_prim = self.decoder[p](concat)
			points_list.append(points_prim)
		points = torch.cat(points_list,dim=-1).permute(0,2,1)
		return points

	def forward(self,opt,image,regular=False):
		code = self.encoder.forward(image)
		points = self.decoder_forward(opt,code,regular=regular)
		return points

	def get_regular_patch_grid(self,opt):
		N = opt.num_meshgrid
		# vertices (UV space)
		U,V = np.meshgrid(range(N+1),range(N+1))
		U = (U.astype(np.float32)/N).reshape([-1])
		V = (V.astype(np.float32)/N).reshape([-1])
		UV = np.stack([U,V],axis=-1)
		UV = torch.tensor(UV,dtype=torch.float32,device=opt.device)
		# facess
		J,I = np.meshgrid(range(N),range(N))
		face_upper = np.stack([I*(N+1)+J,I*(N+1)+J+1,(I+1)*(N+1)+J],axis=-1).reshape([-1,3])
		face_lower = np.stack([I*(N+1)+J+1,(I+1)*(N+1)+J+1,(I+1)*(N+1)+J],axis=-1).reshape([-1,3])
		faces = np.concatenate([face_upper,face_lower],axis=0)
		faces = torch.tensor(faces,dtype=torch.int32,device=opt.device)
		return UV,faces

	def duplicate_faces_original(self,opt):
		faces_list = [self.faces_regular+(opt.num_meshgrid+1)**2*p for p in range(opt.num_prim)]
		self.faces_regular = torch.cat(faces_list,dim=0)
		return self.faces_regular

# ---------- AtlasNet decoder blackbox below ----------

class PointGenCon(torch.nn.Module):
	def __init__(self,code_size):
		self.bottleneck_size = code_size
		super(PointGenCon,self).__init__()
		self.conv1 = torch.nn.Conv1d(self.bottleneck_size,self.bottleneck_size,1)
		self.conv2 = torch.nn.Conv1d(self.bottleneck_size,self.bottleneck_size//2,1)
		self.conv3 = torch.nn.Conv1d(self.bottleneck_size//2,self.bottleneck_size//4,1)
		self.conv4 = torch.nn.Conv1d(self.bottleneck_size//4,3,1)

		self.th = torch.nn.Tanh()
		self.bn1 = torch.nn.BatchNorm1d(self.bottleneck_size)
		self.bn2 = torch.nn.BatchNorm1d(self.bottleneck_size//2)
		self.bn3 = torch.nn.BatchNorm1d(self.bottleneck_size//4)

	def forward(self,x):
		batchsize = x.size()[0]
		x = F.relu(self.bn1(self.conv1(x)))
		x = F.relu(self.bn2(self.conv2(x)))
		x = F.relu(self.bn3(self.conv3(x)))
		x = self.th(self.conv4(x))
		return x

# ---------- ResNet blackbox below ----------

def resnet18(pretrained=False,**kwargs):
	model = ResNet(BasicBlock,[2,2,2,2],**kwargs)
	if pretrained:
		print(util.magenta("loading pretrained encoder..."))
		weight_dict = model_zoo.load_url("https://download.pytorch.org/models/resnet18-5c106cde.pth")
		block_names = list(set([k.split(".")[0] for k in weight_dict.keys()]))
		for b in block_names:
			if b=="fc": continue
			block_weight_dict = {".".join(k.split(".")[1:]): weight_dict[k] for k in weight_dict if k[:len(b)]==b}
			getattr(model,b).load_state_dict(block_weight_dict)
	return model

class BasicBlock(torch.nn.Module):

	expansion = 1

	def __init__(self,inplanes,planes,stride=1,downsample=None):
		super(BasicBlock,self).__init__()
		self.conv1 = torch.nn.Conv2d(inplanes,planes,kernel_size=3,stride=stride,padding=1,bias=False)
		self.bn1 = torch.nn.BatchNorm2d(planes)
		self.relu = torch.nn.ReLU(inplace=True)
		self.conv2 = torch.nn.Conv2d(planes,planes,kernel_size=3,stride=1,padding=1,bias=False)
		self.bn2 = torch.nn.BatchNorm2d(planes)
		self.downsample = downsample
		self.stride = stride

	def forward(self,x):
		residual = x
		out = self.relu(self.bn1(self.conv1(x)))
		out = self.bn2(self.conv2(out))
		if self.downsample is not None:
			residual = self.downsample(x)
		out += residual
		out = self.relu(out)
		return out

class ResNet(torch.nn.Module):

	def __init__(self,block,layers,num_classes=1000):
		self.inplanes = 64
		super(ResNet,self).__init__()
		self.conv1 = torch.nn.Conv2d(3,64,kernel_size=7,stride=2,padding=3,bias=False)
		self.bn1 = torch.nn.BatchNorm2d(64)
		self.relu = torch.nn.ReLU(inplace=True)
		self.maxpool = torch.nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
		self.layer1 = self._make_layer(block,64,layers[0])
		self.layer2 = self._make_layer(block,128,layers[1],stride=2)
		self.layer3 = self._make_layer(block,256,layers[2],stride=2)
		self.layer4 = self._make_layer(block,512,layers[3],stride=2)
		self.avgpool = torch.nn.AvgPool2d(7)
		self.fc = torch.nn.Linear(512*block.expansion,num_classes)
		for m in self.modules():
			if isinstance(m,torch.nn.Conv2d):
				n = m.kernel_size[0]*m.kernel_size[1]*m.out_channels
				m.weight.data.normal_(0,np.sqrt(2./n))
			elif isinstance(m,torch.nn.BatchNorm2d):
				m.weight.data.fill_(1)
				m.bias.data.zero_()

	def _make_layer(self,block,planes,blocks,stride=1):
		downsample = None
		if stride!=1 or self.inplanes!=planes*block.expansion:
			downsample = torch.nn.Sequential(
				torch.nn.Conv2d(self.inplanes,planes*block.expansion,kernel_size=1,stride=stride,bias=False),
				torch.nn.BatchNorm2d(planes*block.expansion),
			)
		layers = []
		layers.append(block(self.inplanes,planes,stride,downsample))
		self.inplanes = planes*block.expansion
		for i in range(1,blocks):
			layers.append(block(self.inplanes,planes))
		return torch.nn.Sequential(*layers)

	def forward(self,x):
		x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
		x = self.layer1(x)
		x = self.layer2(x)
		x = self.layer3(x)
		x = self.layer4(x)
		x = self.avgpool(x)
		x = x.reshape(x.shape[0],-1)
		x = self.fc(x)
		return x

# ---------- chamfer distance blackbox below ----------

class ChamferDistance(torch.autograd.Function):

	@staticmethod
	def forward(ctx,opt,p1,p2):
		batch_size = p1.shape[0]
		num_p1_points = p1.shape[1]
		num_p2_points = p2.shape[1]
		dist1 = torch.zeros(batch_size,num_p1_points,device=opt.device)
		dist2 = torch.zeros(batch_size,num_p2_points,device=opt.device)
		idx1 = torch.zeros(batch_size,num_p1_points,dtype=torch.int32,device=opt.device)
		idx2 = torch.zeros(batch_size,num_p2_points,dtype=torch.int32,device=opt.device)
		p1 = p1.contiguous()
		p2 = p2.contiguous()
		if "cuda" in opt.device:
			chamfer.forward(p1,p2,dist1,dist2,idx1,idx2)
		else:
			raise NotImplementedError("CPU version not implemented")
		ctx.opt = opt
		ctx.save_for_backward(p1,p2,dist1,dist2,idx1,idx2)
		return dist1,dist2

	@staticmethod
	def backward(ctx,grad_dist1,grad_dist2):
		opt = ctx.opt
		p1,p2,dist1,dist2,idx1,idx2 = ctx.saved_tensors
		grad_p1 = torch.zeros_like(p1)
		grad_p2 = torch.zeros_like(p2)
		if "cuda" in opt.device:
			chamfer.backward(p1,p2,grad_p1,grad_p2,grad_dist1,grad_dist2,idx1,idx2)
		else:
			raise NotImplementedError("CPU version not implemented")
		return None,grad_p1,grad_p2