| 101 | |
| 102 | |
| 103 | class LargeScalePortraitVideos(Dataset): |
| 104 | def __init__(self, txt_path, width, height, n_sample_frames, sample_frame_rate, enable_inpaint=True, face_aligner=None): |
| 105 | self.txt_path = txt_path |
| 106 | self.width = width |
| 107 | self.height = height |
| 108 | self.n_sample_frames = n_sample_frames |
| 109 | self.sample_frame_rate = sample_frame_rate |
| 110 | self.enable_inpaint = enable_inpaint |
| 111 | self.video_files = self._read_txt_file_images() |
| 112 | self.face_aligner = face_aligner |
| 113 | |
| 114 | def _read_txt_file_images(self): |
| 115 | with open(self.txt_path, 'r') as file: |
| 116 | lines = file.readlines() |
| 117 | video_files = [] |
| 118 | for line in lines: |
| 119 | video_file = line.strip() |
| 120 | video_files.append(video_file) |
| 121 | return video_files |
| 122 | |
| 123 | def __len__(self): |
| 124 | return len(self.video_files) |
| 125 | |
| 126 | def frame_count(self, frames_path): |
| 127 | files = os.listdir(frames_path) |
| 128 | png_files = [file for file in files if (file.startswith('frame_') and (file.endswith('.png') or file.endswith('.jpg')))] |
| 129 | png_files_count = len(png_files) |
| 130 | return png_files_count |
| 131 | |
| 132 | def find_frames_list(self, frames_path): |
| 133 | files = os.listdir(frames_path) |
| 134 | image_files = [file for file in files if (file.startswith('frame_') and (file.endswith('.png') or file.endswith('.jpg')))] |
| 135 | if image_files[0].startswith('frame_'): |
| 136 | image_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0])) |
| 137 | else: |
| 138 | image_files.sort(key=lambda x: int(x.split('.')[0])) |
| 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 = [] |