| 74 | |
| 75 | |
| 76 | class PreprocessMediaFile: |
| 77 | def __init__(self, config, support_video=False, framerate=None, round_height=16, round_width=16, round_frames=4): |
| 78 | self.config = config |
| 79 | self.video_clip_mode = config.get('video_clip_mode', 'single_beginning') |
| 80 | print(f'using video_clip_mode={self.video_clip_mode}') |
| 81 | self.pil_to_tensor = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]) |
| 82 | self.support_video = support_video |
| 83 | self.framerate = framerate |
| 84 | print(f'using framerate={self.framerate}') |
| 85 | self.round_height = round_height |
| 86 | self.round_width = round_width |
| 87 | self.round_frames = round_frames |
| 88 | if self.support_video: |
| 89 | assert self.framerate |
| 90 | self.tarfile_map = {} |
| 91 | |
| 92 | def __del__(self): |
| 93 | for tar_f in self.tarfile_map.values(): |
| 94 | tar_f.close() |
| 95 | |
| 96 | def __call__(self, spec, mask_filepath, size_bucket=None): |
| 97 | is_video = (Path(spec[1]).suffix in VIDEO_EXTENSIONS) |
| 98 | |
| 99 | if spec[0] is None: |
| 100 | tar_f = None |
| 101 | filepath_or_file = str(spec[1]) |
| 102 | else: |
| 103 | tar_filename = spec[0] |
| 104 | if tar_filename not in self.tarfile_map: |
| 105 | self.tarfile_map[tar_filename] = tarfile.TarFile(tar_filename) |
| 106 | tar_f = self.tarfile_map[tar_filename] |
| 107 | filepath_or_file = tar_f.extractfile(str(spec[1])) |
| 108 | |
| 109 | if is_video: |
| 110 | assert self.support_video |
| 111 | num_frames = 0 |
| 112 | for frame in imageio.v3.imiter(filepath_or_file, fps=self.framerate): |
| 113 | num_frames += 1 |
| 114 | height, width = frame.shape[:2] |
| 115 | video = imageio.v3.imiter(filepath_or_file, fps=self.framerate) |
| 116 | else: |
| 117 | num_frames = 1 |
| 118 | pil_img = Image.open(filepath_or_file) |
| 119 | height, width = pil_img.height, pil_img.width |
| 120 | video = [pil_img] |
| 121 | |
| 122 | if size_bucket is not None: |
| 123 | size_bucket_width, size_bucket_height, size_bucket_frames = size_bucket |
| 124 | else: |
| 125 | size_bucket_width, size_bucket_height, size_bucket_frames = width, height, num_frames |
| 126 | |
| 127 | height_rounded = round_to_nearest_multiple(size_bucket_height, self.round_height) |
| 128 | width_rounded = round_to_nearest_multiple(size_bucket_width, self.round_width) |
| 129 | frames_rounded = round_down_to_multiple(size_bucket_frames - 1, self.round_frames) + 1 |
| 130 | resize_wh = (width_rounded, height_rounded) |
| 131 | |
| 132 | if mask_filepath: |
| 133 | mask_img = Image.open(mask_filepath).convert('RGB') |
no outgoing calls
no test coverage detected