def test_to(self):
V = 20
tex = Textures(
maps=torch.ones((5, 16, 16, 3)),
faces_uvs=torch.randint(size=(5, 10, 3), low=0, high=V),
verts_uvs=torch.ones((5, V, 2)),
)
device = torch.device("cuda:0")
tex = tex.to(device)
self.assertTrue(tex._faces_uvs_padded.device == device)
self.assertTrue(tex._verts_uvs_padded.device == device)
self.assertTrue(tex._maps_padded.device == device)
def make_mesh(verts: torch.tensor,
faces: np.ndarray,
detach: bool,
textures=None) -> Meshes:
device = torch.device("cuda:0")
if detach:
verts = verts.detach()
# Initialize each vertex to be white in color.
if textures is None:
verts_rgb = torch.ones_like(verts)[None] # (1, V, 3)
textures = Textures(verts_rgb=verts_rgb.to(device))
faces = torch.tensor(np.int32(faces), dtype=torch.long).cuda()
# return Meshes(
# verts=[verts.to(device)],
# faces=[faces.to(device)],
# textures=textures
# )
return Meshes(verts=verts.to(device),
faces=faces.to(device).repeat(verts.shape[0], 1, 1),
textures=textures.to(device))
def main():
setup_information()
landmarks_model = _model_init()
train_loader = get_dataset()
for p in landmarks_model.resnet.conv1.parameters():
p.requires_grad = False
for p in landmarks_model.resnet.bn1.parameters():
p.requires_grad = False
for p in landmarks_model.resnet.layer1.parameters():
p.requires_grad = False
#for p in landmarks_model.resnet.layer1.0.parameters(): p.requires_grad=False
params = []
#params = list(net.parameters())
for p in list(landmarks_model.parameters()):
if p.requires_grad == False: continue
params.append(p)
optimizer = optim.RMSprop(params, lr=LEARNING_RATE)
scheduler = MultiStepLR(optimizer, milestones=MILESTONES, gamma=0.2)
landmark_avg_loss = AverageMeter('landmark_avg_loss', ':.4e')
ks_landmark_avg_loss = AverageMeter('ks_landmark_avg_loss', ':.4e')
coeff_avg_loss = AverageMeter('coeff_avg_loss', ':.4e')
pixel_avg_loss = AverageMeter('pixel_avg_loss', ':.4e')
skin_mask_avg_loss = AverageMeter('skin_mask_loss', ':.4e')
avg_loss = AverageMeter('loss', ':.4e')
#normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std=[0.229, 0.224, 0.225])
#trans = transforms.Compose([
# transforms.ToTensor(),
# normalize,
#])
for epoch in range(MAX_EPOCH):
epoch_time = time.time()
scheduler.step()
count = 0
for i_batch, _sample_batched in enumerate(train_loader):
_inputs = _sample_batched[
0] #--- batch_size x 3 x 224 x 224, bgr, 0, 255
_bgr_images = _sample_batched[1].float(
) #--batch_size x 224 x 224 x 3, bgr, [0, 255]
_mask_images = _sample_batched[2].ge(200)
_gt_pts = _sample_batched[3]
_inputs = _inputs / 255.
_inputs = _inputs.sub_(mean[:, None, None]).div_(std[:, None,
None])
_gt_pts = _gt_pts / CROP_SIZE * image_size
_gt_pts = _gt_pts.view(batch_size, -1, 3)
_inputs = _inputs.to(device)
_bgr_images = _bgr_images.to(device)
_mask_images = _mask_images.to(device)
_gt_pts = _gt_pts.to(device)
_render_images = _bgr_images.clone()
coeff = landmarks_model(_inputs)
#-----rendering-----------------------
id_coeff, ex_coeff, tex_coeff, angles, gamma, translation = landmarks_model.Split_coeff(
coeff)
coeff_loss = landmarks_model.get_coeff_loss(
id_coeff, ex_coeff, tex_coeff)
face_shape = landmarks_model.Shape_formation(
id_coeff, ex_coeff, batch_size)
face_norm = landmarks_model.Compute_norm(face_shape, batch_size)
rotation = landmarks_model.Compute_rotation_matrix(angles)
face_shape = torch.matmul(face_shape, rotation) ###旋转vertex
face_shape = face_shape + translation.view(-1, 1, 3).repeat(
1,
face_shape.size()[1], 1)
norm_r = torch.matmul(face_norm, rotation)
face_texture = landmarks_model.Texture_formation(
tex_coeff, batch_size)
face_color, _ = landmarks_model.Illumination_layer(
face_texture, norm_r, gamma)
face_color = Textures(
verts_rgb=face_color.to(device)) #---改成pytorch3d中的Textures数据格式
skin_mask_color = face_texture[:, landmarks_model.skinmask, :]
skin_mask_loss = landmarks_model.get_skin_mask_loss(
skin_mask_color)
mesh = Meshes(face_shape.to(device), landmarks_model.face_index,
face_color)
#---landmarks------
transformed_face_shape = cameras.transform_points(face_shape)
landmarks = transformed_face_shape[:, landmarks_model.facemodel.
keypoints, :]
landmarks = ((landmarks + 1) * image_size - 1) / 2.
landmarks[:, :, :2] = image_size - landmarks[:, :, :
2] #---x坐标和y坐标都需要倒置一下
landmark_loss = landmarks_model.get_landmark_loss(
_gt_pts[:, :, :2], landmarks[:, :, :2])
#-------rendered images---
images = renderer(mesh)
images = images[:, :, :, :3] #---get images
images = images[:, :, :, [2, 1, 0]] #---rgb to bgr
images_clone = images.clone()
#images = images.clamp(0, 255)
index = (images > 0)
_render_images[index] = images[index]
target_images_dir = "debug_images_dir"
if not os.path.exists(target_images_dir):
os.makedirs(target_images_dir)
image_leve_loss = landmarks_model.get_image_level_loss(
_render_images, _bgr_images, _mask_images)
#landmark_loss = 1.6e-3 * landmark_loss
#image_leve_loss = 1.9 * image_leve_loss
#coeff_loss = 3e-4 * coeff_loss
#skin_mask_loss = 5 * skin_mask_loss
landmark_loss = 0.5 * landmark_loss
image_leve_loss = 0.1 * image_leve_loss
coeff_loss = coeff_loss
skin_mask_loss = skin_mask_loss
loss = image_leve_loss + coeff_loss + skin_mask_loss + landmark_loss
#_bgr_images = _bgr_images.cpu().detach().numpy()
#for i in range(batch_size):
# a_image = _bgr_images[i]
# a_target_image_path = os.path.join(target_images_dir, str(i) + '.jpg')
# cv2.imwrite(a_target_image_path, a_image)
#sys.exit(0)
#_render_images = _render_images.cpu().detach().numpy()
#for i in range(batch_size):
# a_image = _render_images[i]
# a_target_image_path = os.path.join(target_images_dir, str(i) + '.jpg')
# cv2.imwrite(a_target_image_path, a_image)
avg_loss.update(loss.detach().item())
landmark_avg_loss.update(landmark_loss.detach().item())
skin_mask_avg_loss.update(skin_mask_loss.detach().item())
coeff_avg_loss.update(coeff_loss.detach().item())
pixel_avg_loss.update(image_leve_loss.detach().item())
if count % 100 == 0:
print('Iter: [%d, %5d]' % (epoch, i_batch))
print(' Iter: [%d, %5d]' % (epoch, i_batch) +
' landmark_loss' + ': %.3e' % landmark_avg_loss.avg)
print(' Iter: [%d, %5d]' % (epoch, i_batch) + ' coeff_loss' +
': %.3e' % coeff_avg_loss.avg)
print(' Iter: [%d, %5d]' % (epoch, i_batch) +
' skin_mask_loss' + ': %.3e' % skin_mask_avg_loss.avg)
print(' Iter: [%d, %5d]' % (epoch, i_batch) +
' image_leve_loss' + ': %.3e' % pixel_avg_loss.avg)
print(' Iter: [%d, %5d]' % (epoch, i_batch) + ' loss' +
': %.3e' % avg_loss.avg)
print('\n')
_render_images = _render_images.cpu().detach().numpy()
_bgr_images = _bgr_images.cpu().detach().numpy()
for i in range(batch_size):
a_image = _render_images[i]
b_image = _bgr_images[i]
c_image = np.concatenate((a_image, b_image), axis=1)
a_target_image_path = os.path.join(target_images_dir,
str(i) + '.jpg')
cv2.imwrite(a_target_image_path, c_image)
landmark_avg_loss = AverageMeter('landmark_avg_loss', ':.4e')
skin_mask_loss = AverageMeter('skin_mask_loss', ':.4e')
pixel_avg_loss = AverageMeter('pixel_avg_loss', ':.4e')
avg_loss = AverageMeter('loss', ':.4e')
if count % 500 == 0:
a_save_name = "_".join([
SUFFIX, 'iter', 'epoch',
'%d' % epoch, 'i_batch',
'%d' % i_batch
]) + '.pth'
a_save_path = os.path.join(WRITE_SNAPSHOT_PATH, a_save_name)
torch.save(landmarks_model.state_dict(), a_save_path)
optimizer.zero_grad()
loss.backward()
optimizer.step()
count += 1
