add inference script features

climategan_wrapper.py

@@ -5,7 +5,7 @@ import os
 import re
 from pathlib import Path
 from uuid import uuid4
-
+from minydra import resolved_args
 import numpy as np
 import torch
 from diffusers import StableDiffusionInpaintPipeline
@@ -541,3 +541,84 @@ class ClimateGAN:
                 im = Image.fromarray(uint8(im))
                 imstem = f"{im_path.stem}---{overwrite_prefix}{painter_prefix}_{event}"
                 im.save(im_path.parent / (imstem + im_path.suffix))
+
+
+if __name__ == "__main__":
+    print("Run `$ python climategan_wrapper.py help` for usage instructions\n")
+
+    # parse arguments
+    args = resolved_args(
+        defaults={
+            "input_folder": None,
+            "output_folder": None,
+            "painter": "both",
+            "help": False,
+        }
+    )
+
+    # print help
+    if args.help:
+        print(
+            "Usage: python inference.py input_folder=/path/to/folder\n"
+            + "By default inferences will be stored in the input folder.\n"
+            + "Add `output_folder=/path/to/folder` for a different output folder.\n"
+            + "By default, both ClimateGAN and Stable Diffusion will be used."
+            + "Change this by adding `painter=climategan` or"
+            + " `painter=stable_diffusion`.\n"
+            + "Make sure you have agreed to the terms of use for the models."
+            + "In particular, visit SD's model card to agree to the terms of use:"
+            + " https://huggingface.co/runwayml/stable-diffusion-inpainting"
+        )
+    # print args
+    args.pretty_print()
+
+    # load models
+    cg = ClimateGAN("models/climategan")
+
+    # check painter type
+    assert args.painter in {"climategan", "stable_diffusion", "both",}, (
+        f"Unknown painter {args.painter}. "
+        + "Allowed values are 'climategan', 'stable_diffusion' and 'both'."
+    )
+
+    # load SD pipeline if need be
+    if args.painter != "climate_gan":
+        cg._setup_stable_diffusion()
+
+    # resolve input folder path
+    in_path = Path(args.input_folder).expanduser().resolve()
+    assert in_path.exists(), f"Folder {str(in_path)} does not exist"
+
+    # output is input if not specified
+    if args.output_folder is None:
+        out_path = in_path
+
+    # find images in input folder
+    im_paths = [
+        p
+        for p in in_path.iterdir()
+        if p.suffix.lower() in [".jpg", ".png", ".jpeg"] and "---" not in p.name
+    ]
+    assert im_paths, f"No images found in {str(im_paths)}"
+
+    print(f"\nFound {len(im_paths)} images in {str(in_path)}\n")
+
+    # infer and write
+    for i, im_path in enumerate(im_paths):
+        print(">>> Processing", f"{i}/{len(im_paths)}", im_path.name)
+        outs = cg.infer_single(
+            np.array(Image.open(im_path)),
+            args.painter,
+            as_pil_image=True,
+            concats=[
+                "input",
+                "masked_input",
+                "climategan_flood",
+                "stable_copy_flood",
+            ],
+        )
+        for k, v in outs.items():
+            name = f"{im_path.stem}---{k}{im_path.suffix}"
+            im = Image.fromarray(uint8(v))
+            im.save(out_path / name)
+        print(">>> Done", f"{i}/{len(im_paths)}", im_path.name, end="\n\n")