(decoded, name)
| 31 | |
| 32 | |
| 33 | def write_image(decoded, name): |
| 34 | assert decoded.ndim == 5 and decoded.shape[2] == 1, decoded.shape |
| 35 | decoded = decoded.squeeze(0) |
| 36 | decoded = ((decoded + 1) / 2).clamp(0, 1) |
| 37 | |
| 38 | img = decoded.squeeze(1) |
| 39 | pil_img = torchvision.transforms.functional.to_pil_image(img) |
| 40 | output_path = args.input.with_name(name + '.jpg') |
| 41 | pil_img.save(output_path) |
| 42 | |
| 43 | |
| 44 | if __name__ == '__main__': |