if not os.path.exists(model_dir):
os.makedirs(model_dir)
# weight_dir = args.weight_dir
mini = not args.complete
full_y = args.full_y
shuffle = not args.no_shuffle
mini = False
full_y = False
shuffle = True
params = get_general_params()
params['IMG_HEIGHT'] = 256
params['IMG_WIDTH'] = 256
params['posemap_downsample'] = 2
ds = mtdataset(params, mini=mini, full_y=full_y)
dl = DataLoader(ds, bs, shuffle)
model = MModel(params, use_cuda=True).cuda()
model.train()
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0424-gan/g_epoch_1670.pth'), strict=False)
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0429-256-gan/g_epoch_720.pth'))
model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0604-gan/epoch_290.pth'))
# if start_epoch != 0:
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/'+weight_dir+'/epoch_%d.pth'%(start_epoch-1)))
vgg_loss = VGGPerceptualLoss().cuda()
l1mask_loss = L1MaskLoss().cuda()
model_params = model.parameters()
optimizer = torch.optim.Adam(model_params, lr=lr)
writer = SummaryWriter('runs/'+model_dir)
print('len(dl)', len(dl))
parser.add_argument('--full_y', action='store_true')
parser.add_argument('--no_shuffle', action='store_true')
args = parser.parse_args()
n_epoch = args.n_epoch
bs = args.bs
start_epoch = args.start_epoch
mini = not args.complete
full_y = args.full_y
shuffle = not args.no_shuffle
args.use_cuda = True
params = get_general_params()
params['IMG_HEIGHT'] = 256
params['IMG_WIDTH'] = 256
params['posemap_downsample'] = 2
GAN = gan(params, args)
ds = mtdataset(params, mini=mini, full_y=full_y)
dl = DataLoader(ds, bs, shuffle)
test_ds = mtdataset(params, mini=False, full_y=full_y, mode='test')
test_dl = DataLoader(test_ds, 16, shuffle=False)
for epoch in range(args.start_epoch + 1, args.n_epoch + args.start_epoch):
if epoch == 1:
GAN.test(test_dl, epoch)
GAN.train(dl, epoch)
if epoch % args.save_freq == 0:
GAN.test(test_dl, epoch)
face_heatmap = torch.from_numpy(face_heatmap)
face_heatmap = torch.sum(face_heatmap, dim=2).unsqueeze(0)
ret[i] = face_heatmap
ret = torch.clamp(ret, 0, 1)
return ret
if __name__ == '__main__':
params = get_general_params()
params['IMG_HEIGHT'] = 256
params['IMG_WIDTH'] = 256
params['posemap_downsample'] = 2
bs = 1
epoch = 0
use_cuda = True
mode = 'test'
ds = mtdataset(params, mini=False, full_y=False, mode=mode)
dl = DataLoader(ds, bs, shuffle=False)
writer = SummaryWriter(log_dir='runs/debug')
dsi = ds[0]
dli = next(iter(dl))
detector = face_detector()
for i, item in enumerate(dsi):
print('ds', i, item.size())
frames = ds.frames
total = 0
good = 0
for i in range(len(ds)):
for j in range(len(frames[i])):
total += 1
src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans, src_mask_gt, tgt_face, tgt_face_box, src_face_box = ds.get(i, [0, j])
# fg_mask = self.FG_mask_Conv(FG)
# fg_mask = F.sigmoid(fg_mask)
# fg_mask = torch.cat((fg_mask, fg_mask, fg_mask), dim=1)
# fg_tgt = torch.mul(fg_tgt, fg_mask)
return fg_tgt, src_mask_delta, src_mask, warped_stack
if __name__ == '__main__':
import cv2
from param import get_general_params
params = get_general_params()
ds = mtdataset(params, full_y=False)
dl = DataLoader(ds, 1)
model = MModel(params, use_cuda=True).cuda()
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/debug')
for src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans, src_mask_gt in dl:
src_img, y, src_pose, tgt_pose, src_mask_prior, x_trans, src_mask_gt = src_img.cuda(
), y.cuda(), src_pose.cuda(), tgt_pose.cuda(), src_mask_prior.cuda(
), x_trans.cuda(), src_mask_gt.cuda()
out = model(src_img, src_pose, tgt_pose, src_mask_prior, x_trans)
# writer.add_images('src_simg', src_img)
# warped = out[-1][0].view(11, 3, out[-1][0].size(1), out[-1][0].size(2))
# print(warped.size(), warped.max(), warped.min(), warped.dtype)
# print(src_img.size(), src_img.max(), src_img.min(), src_img.dtype)