def main():
# model
BASE_MODEL = '../experiment/release_models/'
state_dict = torch.load(BASE_MODEL + 'x2face_model_forpython3.pth')
model = UnwrappedFaceWeightedAverage(output_num_channels=2,
input_num_channels=3,
inner_nc=128)
model.load_state_dict(state_dict['state_dict'])
model = model.cuda()
model = model.eval()
# data
save_path = "results/demo"
# drive_file = ["./examples/Taylor_Swift/1.6/nuBaabkzzzI/"]
# source_files = ["./examples/Taylor_Swift/1.6/nuBaabkzzzI/"]
# source_root = "/home/cxu-serve/p1/common/voxceleb/test/img_sample"
# source_files = []
# for f in os.listdir(source_root):
# source_f = os.path.join(source_root, f)
# if len(os.listdir(source_f)) > 0:
# video_f = os.listdir(source_f)[0]
# img_f = os.listdir(os.path.join(source_f, video_f))[0]
# source_files.append(os.path.join(source_f, video_f, img_f))
drive_file = [
"/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id01567/cIZMA45dX0M/00291_aligned.mp4",
"/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id00017/utfjXffHDgg/00198_aligned.mp4",
"/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id01000/RvjbLfo3XDM/00052_aligned.mp4",
"/home/cxu-serve/p1/common/voxceleb2/unzip/test_video/id04094/2sjuXzB2I1M/00025_aligned.mp4"
]
source_files = drive_file
# drive_file = source_files
# drive_data = Grid(drive_file, nums=10)
# source_data = Grid(source_files, nums=10)
drive_data = Vox(drive_file, nums=10)
source_data = Vox(source_files, nums=8)
drive_loader = DataLoader(dataset=drive_data, batch_size=1, shuffle=False)
source_loader = DataLoader(dataset=source_data,
batch_size=1,
shuffle=False)
# test
with torch.no_grad():
for d_index, drive_imgs in tqdm(drive_loader):
drive_imgs = torch.cat(drive_imgs, dim=0)
drive_imgs = drive_imgs.cuda()
drive = torchvision.utils.make_grid(drive_imgs.cpu().data).permute(
1, 2, 0).numpy()
# source images
for s_index, source_imgs in tqdm(source_loader):
input_imgs = [
img[0].repeat(drive_imgs.shape[0], 1, 1, 1)
for img in source_imgs
]
# get image
result = model(drive_imgs, *input_imgs)
# store
result = result.clamp(min=0, max=1)
result_img = torchvision.utils.make_grid(result.cpu().data)
result_img = result_img.permute(1, 2, 0).numpy()
drive_file = drive_data.get_file(d_index.item()).split('/')[-2]
file_name = os.path.join("{}.{}".format(
*source_data.get_file(s_index.item()).split('/')[-2:]))
file_name = os.path.join(save_path, drive_file, file_name)
if not os.path.exists(file_name):
os.makedirs(file_name)
else:
shutil.rmtree(file_name)
os.makedirs(file_name)
plt.figure()
plt.axis('off')
plt.imshow(np.vstack((result_img, drive)))
plt.savefig(os.path.join(file_name, "result.png"))
plt.close()
source_store = torchvision.utils.make_grid(
torch.cat(source_imgs, dim=0).cpu().data)
source_store = source_store.permute(1, 2, 0).numpy()
plt.figure()
plt.axis('off')
plt.imshow(source_store)
plt.savefig(os.path.join(file_name, "origin.png"))
plt.close()
if opt.copy_weights:
checkpoint_file = torch.load(opt.old_model)
model.load_state_dict(checkpoint_file['state_dict'])
opt.model_epoch_path = opt.model_epoch_path + 'copyWeights'
del checkpoint_file
criterion = L1Loss()
if opt.num_views > 2:
opt.model_epoch_path = opt.model_epoch_path + 'num_views' + str(opt.num_views) + 'combination_function' + opt.combination_function
model.stats = {'photometric_error' : np.zeros((0,1)), 'eyemoutherror' : np.zeros((0,1)), 'contenterror' : np.zeros((0,1)), 'loss' : np.zeros((0,1))}
model.val_stats = {'photometric_error' : np.zeros((0,1)), 'eyemoutherror' : np.zeros((0,1)), 'contenterror' : np.zeros((0,1)), 'loss' : np.zeros((0,1))}
model = model.cuda()
criterion = criterion.cuda()
# optimizer = optim.SGD([{'params' : model.pix2pixUnwrapped.parameters()},
# {'params' : model.pix2pixSampler.parameters(), 'lr' : opt.lr}], lr=opt.lr, momentum=0.9)
parameters = [{'params' : model.parameters()}]
optimizer = optim.SGD(parameters, lr=opt.lr, momentum=0.9)
def run_batch(imgs, requires_grad=False, volatile=False, other_images=None):
for i in range(0, len(imgs)):
imgs[i] = Variable(imgs[i], requires_grad=requires_grad, volatile=volatile).cuda()
poses = imgs[-3]
print('Poses', poses.size())
def main():
# model
BASE_MODEL = '../experiment/release_models/'
state_dict = torch.load(BASE_MODEL + 'x2face_model_forpython3.pth')
model = UnwrappedFaceWeightedAverage(output_num_channels=2,
input_num_channels=3,
inner_nc=128)
model.load_state_dict(state_dict['state_dict'])
model = model.cuda()
model = model.eval()
# data
save_path = "extra_degree_result/vox"
ref_files = []
files = []
pickle_files = np.load('vox_demo.npy', allow_pickle=True)
files = list(set([f[0] for f in pickle_files]))
tgt_ids = {}
for f in pickle_files:
if f[0] not in tgt_ids:
tgt_ids[f[0]] = []
tgt_ids[f[0]] += np.arange(int(f[1]), int(f[2])).tolist()
ref_files = files
# files = files[670:]
total_f = len(files)
files = [f for f in files if os.path.exists(f)]
print('totally {} files while valid {} files'.format(total_f, len(files)))
dataset = VoxSingle(files, ref_indx='0,5,10,15,20,25,30,35')
dataloader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=0)
# test
state = 1
with torch.no_grad():
while state != -1:
print("current {}/{}".format(state, len(files)))
# get reference images
ref_imgs = dataset.get_ref()
# ref_imgs = ref_dataset.get_ref()
# save reference images
file = dataset.cur_file.split('/')
path_name = '{}'.format(file[-1])
save_file_name = os.path.join(save_path, path_name)
for ref_id, img in enumerate(ref_imgs):
save_ref_file = os.path.join(save_file_name, 'reference')
if not os.path.exists(save_ref_file):
os.makedirs(save_ref_file)
save_img = (img * 255).permute(1, 2, 0).cpu().numpy().astype(
np.uint8)
Image.fromarray(save_img).save(
os.path.join(save_ref_file, '%05d.jpg' % ref_id))
# preprocess
ref_imgs = [ref.unsqueeze(0) for ref in ref_imgs]
ref_imgs = torch.cat(ref_imgs, dim=0)
ref_imgs = ref_imgs.cuda()
ref = torchvision.utils.make_grid(ref_imgs.cpu().data).permute(
1, 2, 0).numpy()
# synthesize
for d_index, drive_img in tqdm(dataloader):
if d_index not in tgt_ids[dataset.cur_file]:
continue
if os.path.exists(
os.path.join(save_file_name, '%05d.jpg' % d_index)):
continue
drive_img = drive_img.cuda()
input_imgs = [
img.repeat(drive_img.shape[0], 1, 1, 1) for img in ref_imgs
]
# get image
result = model(drive_img, *input_imgs)
# store
result = result.clamp(min=0, max=1)
result_img = torchvision.utils.make_grid(result.cpu().data)
result_img = result_img.permute(1, 2, 0).numpy()
save_img = drive_img.clamp(min=0, max=1)
save_img = torchvision.utils.make_grid(save_img.cpu().data)
save_img = save_img.permute(1, 2, 0).numpy()
# final_img = np.hstack([result_img, save_img])
final_img = result_img
final_img = (final_img * 255).astype(np.uint8)
Image.fromarray(final_img).save(
os.path.join(save_file_name, '%05d.jpg' % d_index))
# combine video
# image_to_video(save_file_name, os.path.join(save_file_name, '{}.mp4'.format(path_name)))
# new file
state = dataset.nextfile()