feat(doc): update inference models and examples

jolibrain · Sep 1, 2023 · 3c43a7b · 3c43a7b
1 parent 98cef11
commit 3c43a7b
Show file tree

Hide file tree

Showing 11 changed files with 18 additions and 8 deletions.
diff --git a/docs/source/_static/gan_glasses2noglasses_output.jpg b/docs/source/_static/gan_glasses2noglasses_output.jpg
diff --git a/docs/source/_static/noglasses2glasses_ddpm_cond.png b/docs/source/_static/noglasses2glasses_ddpm_cond.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_generated.png b/docs/source/_static/noglasses2glasses_ddpm_generated.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_generated_crop.png b/docs/source/_static/noglasses2glasses_ddpm_generated_crop.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_mask.png b/docs/source/_static/noglasses2glasses_ddpm_mask.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_orig.png b/docs/source/_static/noglasses2glasses_ddpm_orig.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_y_0.png b/docs/source/_static/noglasses2glasses_ddpm_y_0.png
diff --git a/docs/source/_static/noglasses2glasses_ddpm_y_t.png b/docs/source/_static/noglasses2glasses_ddpm_y_t.png
diff --git a/docs/source/inference.rst b/docs/source/inference.rst
@@ -37,11 +37,11 @@ Download a pretrained model:
 
 .. code:: bash
 
-   wget https://joligen.com/models/glasses2noglasses.zip
-   unzip glasses2noglasses.zip
+   wget https://joligen.com/models/joligen_model_gan_glasses2noglasses.zip
+   unzip joligen_model_gan_glasses2noglasses.zip
    mkdir checkpoints
    mv glasses2noglasses/ ./checkpoints/
-   rm glasses2noglasses.zip
+   rm joligen_model_gan_glasses2noglasses.zip
 
 Run the inference script
 ========================
@@ -57,6 +57,7 @@ The output file is the ``target.jpg`` image in the current directory:
    :class: borderless
 
    -  -  .. image:: _static/gan_glasses2noglasses_orig_image.jpg
+           :width: 128
       -  .. image:: _static/gan_glasses2noglasses_output.jpg
 
    -  -  original image given as input to the model
@@ -89,11 +90,11 @@ Download a pretrained model:
 
 .. code:: bash
 
-   wget https://joligen.com/models/noglasses2glasses.zip
-   unzip noglasses2glasses.zip
+   wget https://joligen.com/models/joligen_model_ddpm_noglasses2glasses.zip
+   unzip joligen_model_ddpm_noglasses2glasses.zip
    mkdir checkpoints
    mv noglasses2glasses/ ./checkpoints/
-   rm noglasses2glasses.zip
+   rm joligen_model_ddpm_noglasses2glasses.zip
 
 Run the inference script
 ========================
@@ -102,7 +103,7 @@ Run the inference script
 
    mkdir noglasses2glasses_inference_output
    cd scripts/
-   python3 gen_single_image_diffusion.py --model-in-file ../checkpoints/noglasses2glasses/latest_net_G_A.pth --img-in ../datasets/noglasses2glasses_ffhq/trainA/img/00001.jpg --mask-in ../datasets/noglasses2glasses_ffhq/trainA/bbox/00001.jpg --dir-out ../noglasses2glasses_inference_output --img-width 128 --img-height 128
+   python3 gen_single_image_diffusion.py --model-in-file ../checkpoints/noglasses2glasses/latest_net_G_A.pth --img-in ../datasets/noglasses2glasses_ffhq/trainA/img/00002.jpg --mask-in ../datasets/noglasses2glasses_ffhq/trainA/bbox/00002.jpg --dir-out ../noglasses2glasses_inference_output --img-width 128 --img-height 128
 
 The output files will be in the ``noglasses2glasses_inference_output``
 folder, with:

diff --git a/examples/example_ddpm_noglasses2glasses.json b/examples/example_ddpm_noglasses2glasses.json
@@ -278,7 +278,7 @@
     "checkpoints_dir": "./checkpoints/",
     "dataroot": "noglasses2glasses_ffhq",
     "ddp_port": "12355",
-    "gpu_ids": "1",
+    "gpu_ids": "0",
     "model_type": "palette",
     "name": "noglasses2glasses",
     "phase": "train",

diff --git a/scripts/gen_single_image.py b/scripts/gen_single_image.py
@@ -52,6 +52,12 @@ def load_model(modelpath, model_in_file, cpu, gpuid):
 
 parser.add_argument("--img-in", help="image to transform", required=True)
 parser.add_argument("--img-out", help="transformed image", required=True)
+parser.add_argument(
+    "--img-width", type=int, help="image width, defaults to model crop size"
+)
+parser.add_argument(
+    "--img-height", type=int, help="image height, defaults to model crop size"
+)
 parser.add_argument("--cpu", action="store_true", help="whether to use CPU")
 parser.add_argument("--gpuid", type=int, default=0, help="which GPU to use")
 args = parser.parse_args()
@@ -65,8 +71,11 @@ def load_model(modelpath, model_in_file, cpu, gpuid):
 )
 
 # reading image
+img_width = args.img_width if args.img_width is not None else opt.data_crop_size
+img_height = args.img_height if args.img_height is not None else opt.data_crop_size
 img = cv2.imread(args.img_in)
 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+img = cv2.resize(img, (img_width, img_height), interpolation=cv2.INTER_CUBIC)
 
 # preprocessing
 tranlist = [