(self, idx)
| 139 | return image_files |
| 140 | |
| 141 | def __getitem__(self, idx): |
| 142 | warnings.filterwarnings('ignore', category=DeprecationWarning) |
| 143 | warnings.filterwarnings('ignore', category=FutureWarning) |
| 144 | |
| 145 | frames_path = os.path.join(self.video_files[idx], "images") |
| 146 | face_masks_path = os.path.join(self.video_files[idx], "face_masks") |
| 147 | lip_masks_path = os.path.join(self.video_files[idx], "lip_masks") |
| 148 | |
| 149 | video_length = self.frame_count(frames_path) |
| 150 | frames_list = self.find_frames_list(frames_path) |
| 151 | all_indices = list(range(0, video_length)) |
| 152 | clip_length = min(video_length, (self.n_sample_frames - 1) * self.sample_frame_rate + 1) |
| 153 | |
| 154 | start_idx = random.randint(0, video_length - clip_length) |
| 155 | batch_index = np.linspace(start_idx, start_idx + clip_length - 1, self.n_sample_frames, dtype=int).tolist() |
| 156 | |
| 157 | tgt_pil_image_list = [] |
| 158 | tgt_bgr_image_list = [] |
| 159 | tgt_face_masks_list = [] |
| 160 | tgt_lip_masks_list = [] |
| 161 | |
| 162 | reference_frame_path = os.path.join(frames_path, frames_list[start_idx]) |
| 163 | reference_pil_image = Image.open(reference_frame_path).convert('RGB') |
| 164 | reference_pil_image = reference_pil_image.resize([self.width, self.height], Image.LANCZOS) |
| 165 | reference_pil_image = torch.from_numpy(np.array(reference_pil_image)).float() |
| 166 | reference_pil_image = reference_pil_image / 127.5 - 1 |
| 167 | |
| 168 | for index in batch_index: |
| 169 | tgt_img_path = os.path.join(frames_path, frames_list[index]) |
| 170 | file_name = os.path.basename(tgt_img_path) |
| 171 | face_mask_path = os.path.join(face_masks_path, file_name) |
| 172 | lip_mask_path = os.path.join(lip_masks_path, file_name) |
| 173 | try: |
| 174 | tgt_img_pil = Image.open(tgt_img_path).convert('RGB') |
| 175 | tgt_img_pil = tgt_img_pil.resize([self.width, self.height], Image.LANCZOS) |
| 176 | tgt_img_tensor = torch.from_numpy(np.array(tgt_img_pil)).float() |
| 177 | tgt_img_normalized = tgt_img_tensor / 127.5 - 1 |
| 178 | tgt_pil_image_list.append(tgt_img_normalized) |
| 179 | except Exception as e: |
| 180 | print(f"Fail loading the image: {tgt_img_path}") |
| 181 | print(1/0) |
| 182 | |
| 183 | try: |
| 184 | bgr_img = cv2.imread(tgt_img_path) |
| 185 | tgt_bgr_image_list.append(bgr_img.copy()) |
| 186 | except Exception as e: |
| 187 | print(f"Fail loading the bgr image: {tgt_img_path}") |
| 188 | print(1/0) |
| 189 | |
| 190 | try: |
| 191 | tgt_lip_mask = Image.open(lip_mask_path) |
| 192 | tgt_lip_mask = tgt_lip_mask.resize([self.width, self.height], Image.LANCZOS) |
| 193 | tgt_lip_mask = torch.from_numpy(np.array(tgt_lip_mask)).float() |
| 194 | tgt_lip_mask = tgt_lip_mask / 255 |
| 195 | except Exception as e: |
| 196 | print(f"Fail loading the lip masks: {lip_mask_path}") |
| 197 | tgt_lip_mask = torch.ones(self.height, self.width) |
| 198 | tgt_lip_masks_list.append(tgt_lip_mask) |
nothing calls this directly
no test coverage detected