Merge pull request #36 from jmisilo/develop

Develop

.gitignore

@@ -2,6 +2,6 @@
 /data/
 __pycache__/
 .vscode/
-/weights/arch/
+/weights/
 /wandb/
 test_generation.ipynb

src/evaluate.py

@@ -15,7 +15,7 @@
 from model.loops import evaluate_dataset
 from model.model import Net
 from utils.config import Config
-from utils.load_ckp import load_ckp
+from utils.load_ckp import download_weights
 
 config = Config()
 parser = argparse.ArgumentParser()
@@ -48,9 +48,7 @@
 
 ckp_path = os.path.join(config.weights_dir, args.checkpoint_name)
 
-assert os.path.isfile(ckp_path), 'Checkpoint does not exist'
 assert os.path.exists(args.img_path), 'Path to the test image folder does not exist'
-assert os.path.exists(args.res_path), 'Path to the results folder does not exist'
 
 # set seed
 random.seed(config.seed)
@@ -81,7 +79,14 @@
 
     _, _, test_dataset = random_split(dataset, [config.train_size, config.val_size, config.test_size])
 
-    load_ckp(ckp_path, model, device=device)
+    if not os.path.exists(config.weights_dir):
+        os.makedirs(config.weights_dir)
+
+    if not os.path.isfile(ckp_path):
+        download_weights(ckp_path)
+
+    checkpoint = torch.load(ckp_path, map_location=device)
+    model.load_state_dict(checkpoint)    
 
     save_path = os.path.join(args.res_path, args.checkpoint_name[:-3])
 

src/predict.py

@@ -13,8 +13,8 @@
 import torch
 
 from model.model import Net
+from utils.load_ckp import download_weights
 from utils.config import Config
-from utils.load_ckp import load_ckp
 
 config = Config()
 parser = argparse.ArgumentParser()
@@ -59,7 +59,6 @@
     ckp_path = os.path.join(config.weights_dir, args.checkpoint_name)
 
     assert os.path.isfile(args.img_path), 'Image does not exist'
-    assert os.path.isfile(ckp_path), 'Checkpoint does not exist'
 
     if not os.path.exists(args.res_path):
         os.makedirs(args.res_path)
@@ -75,8 +74,15 @@
         max_len=config.max_len,
         device=device
     )
-
-    load_ckp(ckp_path, model, device=device)
+
+    if not os.path.exists(config.weights_dir):
+        os.makedirs(config.weights_dir)
+
+    if not os.path.isfile(ckp_path):
+        download_weights(ckp_path)
+
+    checkpoint = torch.load(ckp_path, map_location=device)
+    model.load_state_dict(checkpoint)    
 
     model.eval()
 

src/utils/load_ckp.py

@@ -1,7 +1,8 @@
 '''
-    Utility function to load checkpoint from coresponding file.
+    Utility functions for loading weights.
 '''
 
+import gdown
 import torch
 
 def load_ckp(checkpoint_fpath, model, optimizer=None, scheduler=None, scaler=None, device='cpu'):
@@ -21,4 +22,11 @@ def load_ckp(checkpoint_fpath, model, optimizer=None, scheduler=None, scaler=Non
     if scaler is not None:
         scaler.load_state_dict(checkpoint['scaler_state_dict'])
 
-    return checkpoint['epoch']
+    return checkpoint['epoch']
+
+def download_weights(checkpoint_fpath):
+    '''
+        Downloads weights from Google Drive.
+    '''
+
+    gdown.download('https://drive.google.com/uc?id=1lEufQVOETFEIhPdFDYaez31uroq_5Lby', checkpoint_fpath, quiet=False)