(
self,
base_path=None, metadata_path=None,
num_frames=81,
time_division_factor=4, time_division_remainder=1,
max_pixels=1920*1080, height=None, width=None,
height_division_factor=16, width_division_factor=16,
data_file_keys=("video",),
image_file_extension=("jpg", "jpeg", "png", "webp"),
video_file_extension=("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm", "gif"),
repeat=1,
args=None,
)
| 148 | |
| 149 | class VideoDataset(torch.utils.data.Dataset): |
| 150 | def __init__( |
| 151 | self, |
| 152 | base_path=None, metadata_path=None, |
| 153 | num_frames=81, |
| 154 | time_division_factor=4, time_division_remainder=1, |
| 155 | max_pixels=1920*1080, height=None, width=None, |
| 156 | height_division_factor=16, width_division_factor=16, |
| 157 | data_file_keys=("video",), |
| 158 | image_file_extension=("jpg", "jpeg", "png", "webp"), |
| 159 | video_file_extension=("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm", "gif"), |
| 160 | repeat=1, |
| 161 | args=None, |
| 162 | ): |
| 163 | if args is not None: |
| 164 | base_path = args.dataset_base_path |
| 165 | metadata_path = args.dataset_metadata_path |
| 166 | height = args.height |
| 167 | width = args.width |
| 168 | max_pixels = args.max_pixels |
| 169 | num_frames = args.num_frames |
| 170 | data_file_keys = args.data_file_keys.split(",") |
| 171 | repeat = args.dataset_repeat |
| 172 | |
| 173 | self.base_path = base_path |
| 174 | self.num_frames = num_frames |
| 175 | self.time_division_factor = time_division_factor |
| 176 | self.time_division_remainder = time_division_remainder |
| 177 | self.max_pixels = max_pixels |
| 178 | self.height = height |
| 179 | self.width = width |
| 180 | self.height_division_factor = height_division_factor |
| 181 | self.width_division_factor = width_division_factor |
| 182 | self.data_file_keys = data_file_keys |
| 183 | self.image_file_extension = image_file_extension |
| 184 | self.video_file_extension = video_file_extension |
| 185 | self.repeat = repeat |
| 186 | |
| 187 | if height is not None and width is not None: |
| 188 | print("Height and width are fixed. Setting `dynamic_resolution` to False.") |
| 189 | self.dynamic_resolution = False |
| 190 | elif height is None and width is None: |
| 191 | print("Height and width are none. Setting `dynamic_resolution` to True.") |
| 192 | self.dynamic_resolution = True |
| 193 | |
| 194 | if metadata_path is None: |
| 195 | print("No metadata. Trying to generate it.") |
| 196 | metadata = self.generate_metadata(base_path) |
| 197 | print(f"{len(metadata)} lines in metadata.") |
| 198 | self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] |
| 199 | elif metadata_path.endswith(".json"): |
| 200 | with open(metadata_path, "r") as f: |
| 201 | metadata = json.load(f) |
| 202 | self.data = metadata |
| 203 | else: |
| 204 | metadata = pd.read_csv(metadata_path) |
| 205 | self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] |
| 206 | |
| 207 |
nothing calls this directly
no test coverage detected