(self, spec, mask_filepath, size_bucket=None)
| 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') |
| 134 | img_hw = (height, width) |
| 135 | mask_hw = (mask_img.height, mask_img.width) |
| 136 | if mask_hw != img_hw: |
| 137 | raise ValueError( |
| 138 | f'Mask shape {mask_hw} was not the same as image shape {img_hw}.\n' |
| 139 | f'Image path: {spec[1]}\n' |
| 140 | f'Mask path: {mask_filepath}' |
| 141 | ) |
| 142 | mask_img = ImageOps.fit(mask_img, resize_wh) |
| 143 | mask = torchvision.transforms.functional.to_tensor(mask_img)[0].to(torch.float16) # use first channel |
| 144 | else: |
| 145 | mask = None |
| 146 | |
| 147 | resized_video = torch.empty((num_frames, 3, height_rounded, width_rounded)) |
| 148 | for i, frame in enumerate(video): |
| 149 | if not isinstance(frame, Image.Image): |
| 150 | frame = torchvision.transforms.functional.to_pil_image(frame) |
| 151 | cropped_image = convert_crop_and_resize(frame, resize_wh) |
| 152 | resized_video[i, ...] = self.pil_to_tensor(cropped_image) |
| 153 |
nothing calls this directly
no test coverage detected