Correcting issues for first time setup

SixRoses666 · Mar 7, 2022 · 4e05cda · 4e05cda
1 parent 2f136fc
commit 4e05cda
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -23,4 +23,13 @@ cd src
 python download_models.py
 ```
 
+#### Download Datasets
+
+First, change the `pytorch_data_dir` variable to your 
+systems pytorch data directory where datasets are stored. 
+
+```
+python download_data.py
+```
+
 
diff --git a/src/download_models.py b/src/download_models.py
@@ -1,13 +1,16 @@
 from os.path import join, exists
 import wget
+import os
 
 models_dir = join("..", "models")
+os.makedirs(models_dir, exist_ok=True)
 model_url_root = "https://marhamilresearch4.blob.core.windows.net/stego-public/models/models/"
 model_names = ["moco_v2_800ep_pretrain.pth.tar",
                "model_epoch_0720_iter_085000.pth",
                "picie.pkl"]
 
 saved_models_dir = join("..", "saved_models")
+os.makedirs(saved_models_dir, exist_ok=True)
 saved_model_url_root = "https://marhamilresearch4.blob.core.windows.net/stego-public/saved_models/"
 saved_model_names = ["cityscapes_vit_base_1.ckpt",
                      "cocostuff27_vit_base_5.ckpt",