Spaces:
Runtime error
Runtime error
| """This script is the test script for Deep3DFaceRecon_pytorch | |
| """ | |
| import os | |
| import numpy as np | |
| import torch | |
| from data import create_dataset | |
| from data.flist_dataset import default_flist_reader | |
| from options.test_options import TestOptions | |
| from PIL import Image | |
| from scipy.io import loadmat | |
| from scipy.io import savemat | |
| from util.load_mats import load_lm3d | |
| from util.preprocess import align_img | |
| from util.visualizer import MyVisualizer | |
| from models import create_model | |
| def get_data_path(root="examples"): | |
| im_path = [os.path.join(root, i) for i in sorted(os.listdir(root)) if i.endswith("png") or i.endswith("jpg")] | |
| lm_path = [i.replace("png", "txt").replace("jpg", "txt") for i in im_path] | |
| lm_path = [ | |
| os.path.join(i.replace(i.split(os.path.sep)[-1], ""), "detections", i.split(os.path.sep)[-1]) for i in lm_path | |
| ] | |
| return im_path, lm_path | |
| def read_data(im_path, lm_path, lm3d_std, to_tensor=True): | |
| # to RGB | |
| im = Image.open(im_path).convert("RGB") | |
| W, H = im.size | |
| lm = np.loadtxt(lm_path).astype(np.float32) | |
| lm = lm.reshape([-1, 2]) | |
| lm[:, -1] = H - 1 - lm[:, -1] | |
| _, im, lm, _ = align_img(im, lm, lm3d_std) | |
| if to_tensor: | |
| im = torch.tensor(np.array(im) / 255.0, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) | |
| lm = torch.tensor(lm).unsqueeze(0) | |
| return im, lm | |
| def main(rank, opt, name="examples"): | |
| device = torch.device(rank) | |
| torch.cuda.set_device(device) | |
| model = create_model(opt) | |
| model.setup(opt) | |
| model.device = device | |
| model.parallelize() | |
| model.eval() | |
| visualizer = MyVisualizer(opt) | |
| im_path, lm_path = get_data_path(name) | |
| lm3d_std = load_lm3d(opt.bfm_folder) | |
| for i in range(len(im_path)): | |
| print(i, im_path[i]) | |
| img_name = im_path[i].split(os.path.sep)[-1].replace(".png", "").replace(".jpg", "") | |
| if not os.path.isfile(lm_path[i]): | |
| print("%s is not found !!!" % lm_path[i]) | |
| continue | |
| im_tensor, lm_tensor = read_data(im_path[i], lm_path[i], lm3d_std) | |
| data = {"imgs": im_tensor, "lms": lm_tensor} | |
| model.set_input(data) # unpack data from data loader | |
| model.test() # run inference | |
| visuals = model.get_current_visuals() # get image results | |
| visualizer.display_current_results( | |
| visuals, | |
| 0, | |
| opt.epoch, | |
| dataset=name.split(os.path.sep)[-1], | |
| save_results=True, | |
| count=i, | |
| name=img_name, | |
| add_image=False, | |
| ) | |
| model.save_mesh( | |
| os.path.join( | |
| visualizer.img_dir, name.split(os.path.sep)[-1], "epoch_%s_%06d" % (opt.epoch, 0), img_name + ".obj" | |
| ) | |
| ) # save reconstruction meshes | |
| model.save_coeff( | |
| os.path.join( | |
| visualizer.img_dir, name.split(os.path.sep)[-1], "epoch_%s_%06d" % (opt.epoch, 0), img_name + ".mat" | |
| ) | |
| ) # save predicted coefficients | |
| if __name__ == "__main__": | |
| opt = TestOptions().parse() # get test options | |
| main(0, opt, opt.img_folder) | |