| 133 | os.makedirs(self.save_dir, exist_ok=True) |
| 134 | |
| 135 | def load_lq(self) -> Generator[Image.Image, None, None]: |
| 136 | img_exts = [".png", ".jpg", ".jpeg"] |
| 137 | assert os.path.isdir( |
| 138 | self.args.input |
| 139 | ), "Please put your low-quality images in a folder." |
| 140 | for file_name in sorted(os.listdir(self.args.input)): |
| 141 | stem, ext = os.path.splitext(file_name) |
| 142 | if ext not in img_exts: |
| 143 | print(f"{file_name} is not an image, continue") |
| 144 | continue |
| 145 | file_path = os.path.join(self.args.input, file_name) |
| 146 | lq = Image.open(file_path).convert("RGB") |
| 147 | print(f"load lq: {file_path}") |
| 148 | self.loop_ctx["file_stem"] = stem |
| 149 | yield lq |
| 150 | |
| 151 | def after_load_lq(self, lq: Image.Image) -> np.ndarray: |
| 152 | return np.array(lq) |