(self, video_path=None, image_path=None)
| 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): |
| 118 | src_image_pil_256 = src_image_pil.resize((256, 256)) |
| 119 | os.makedirs(os.path.join(os.path.dirname(image_path), '3dmm'), exist_ok=True) |
| 120 | lm_np = get_landmark([src_image_pil_256]) |
| 121 | # print(lm_np.shape) |
| 122 | source_3dmm = self.model_3dmm.get_3dmm([src_image_pil_256], lm_np) |
| 123 | # print(coeff_3dmm.shape) |
| 124 | np.save(source_3dmm_path, source_3dmm) |
| 125 | |
| 126 | source_3dmm = np.load(source_3dmm_path, allow_pickle=True) |
| 127 | source_3dmm = torch.from_numpy(source_3dmm) |
| 128 | else: |
| 129 | src_image_pil = frames_pil[0] |
| 130 | source_3dmm = coeff_3dmm[0].unsqueeze(0) |
| 131 | |
| 132 | if self.if_align: |
| 133 | src_lm_np = get_landmark([src_image_pil]) |
| 134 | src_align_pil = align_image_pil([src_image_pil], src_lm_np) |
| 135 | src_align_pil = src_align_pil[0] |
| 136 | else: |
| 137 | src_align_pil = src_image_pil |
| 138 | |
| 139 | return { |
| 140 | 'source_align': src_align_pil, |
| 141 | 'source_image': src_image_pil, |
| 142 | 'source_3dmm': source_3dmm, |
| 143 | 'frames': frames_pil if video_path is not None else None, |
| 144 | 'coeff_3dmm': coeff_3dmm if video_path is not None else None, |
| 145 | 'video_name': video_name if video_path is not None else None |
| 146 | } |
| 147 | |
| 148 | def transform_semantic(self, semantic, frame_index): |
no test coverage detected