| 58 | |
| 59 | |
| 60 | class TempVideoDataset: |
| 61 | |
| 62 | def __init__(self, video_list, model_3dmm, if_align=False, cross_id=False, image_list=None, resize=256): |
| 63 | self.video_list = video_list |
| 64 | self.model_3dmm = model_3dmm |
| 65 | self.cross_id = cross_id |
| 66 | self.image_list = image_list |
| 67 | self.if_align = if_align |
| 68 | if self.cross_id and len(self.video_list) != len(self.image_list) and len(self.video_list) > 0: |
| 69 | self.video_list = self.video_list * (len(self.image_list) // len(self.video_list) + 1) |
| 70 | self.video_list = self.video_list[:len(self.image_list)] |
| 71 | |
| 72 | self.video_index = -1 |
| 73 | self.transform = transforms.Compose([ |
| 74 | transforms.Resize((resize, resize)), |
| 75 | transforms.ToTensor(), |
| 76 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), |
| 77 | ]) |
| 78 | self.image_transform = transforms.Compose([ |
| 79 | transforms.ToTensor(), |
| 80 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), |
| 81 | ]) |
| 82 | self.semantic_radius = 13 |
| 83 | self.croper = Croper() |
| 84 | |
| 85 | def __len__(self): |
| 86 | if len(self.video_list) > 0: |
| 87 | return len(self.video_list) |
| 88 | else: |
| 89 | return len(self.image_list) |
| 90 | |
| 91 | def data_preprocess(self, video_path=None, image_path=None): |
| 92 | # Hard code; Bad writing |
| 93 | if video_path is not None: |
| 94 | video_name = os.path.basename(video_path).split('.')[0] |
| 95 | frames_pil = video_util.read_video(video_path, resize=256) |
| 96 | |
| 97 | save_3dmm_path = os.path.join(os.path.dirname(video_path), '3dmm', '3dmm_' + video_name + '.npy') |
| 98 | if not os.path.exists(save_3dmm_path): |
| 99 | os.makedirs(os.path.join(os.path.dirname(video_path), '3dmm'), exist_ok=True) |
| 100 | lm_np = get_landmark(frames_pil) |
| 101 | |
| 102 | frames_pil = self.croper.crop(frames_pil, lm_np) |
| 103 | lm_np = get_landmark(frames_pil) |
| 104 | |
| 105 | coeff_3dmm = self.model_3dmm.get_3dmm(frames_pil, lm_np) |
| 106 | # print(coeff_3dmm.shape) |
| 107 | np.save(save_3dmm_path, coeff_3dmm) |
| 108 | |
| 109 | coeff_3dmm = np.load(save_3dmm_path, allow_pickle=True) |
| 110 | coeff_3dmm = torch.from_numpy(coeff_3dmm) |
| 111 | |
| 112 | if self.cross_id and image_path is not None: |
| 113 | src_image_pil = Image.open(image_path).convert("RGB") # prevent png exist channel error |
| 114 | |
| 115 | image_name = os.path.basename(image_path).split('.')[0] |
| 116 | source_3dmm_path = os.path.join(os.path.dirname(image_path), '3dmm', '3dmm_' + image_name + '.npy') |
| 117 | if not os.path.exists(source_3dmm_path): |
no outgoing calls
no test coverage detected