Modifying App.py

app.py

@@ -19,6 +19,26 @@ def cropping(img):
 
     return img
 
+def load_model(ckpt, model, optimizer=None):
+    ckpt_dict = torch.load(ckpt, map_location='cpu')
+    # keep backward compatibility
+    if 'model' not in ckpt_dict and 'optimizer' not in ckpt_dict:
+        state_dict = ckpt_dict
+    else:
+        state_dict = ckpt_dict['model']
+    weights = {}
+    for key, value in state_dict.items():
+        if key.startswith('module.'):
+            weights[key[len('module.'):]] = value
+        else:
+            weights[key] = value
+
+    model.load_state_dict(weights)
+
+    if optimizer is not None:
+        optimizer_state = ckpt_dict['optimizer']
+        optimizer.load_state_dict(optimizer_state)
+
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(DEVICE)
 CWD = "."
@@ -36,48 +56,35 @@ CKPT_FILE_NAMES = {
 }
 MODEL_CLASSES = {
     'Indoor': {
-        'Resnet_enc':enc_dec_model,
-        'Unet':ResNet18UNet,
-        'Densenet_enc':Densenet
+        'Resnet_enc':enc_dec_model(max_depth = 10),
+        'Unet':ResNet18UNet(max_depth = 10),
+        'Densenet_enc':Densenet(max_depth = 10)
     },
 
     'Outdoor': {
-        'Resnet_enc':enc_dec_model,
-        'Unet':UNetWithResnet50Encoder,
-        'Densenet_enc':Densenet
+        'Resnet_enc':enc_dec_model(max_depth = 80),
+        'Unet':UNetWithResnet50Encoder(max_depth = 80),
+        'Densenet_enc':Densenet(max_depth = 80)
     },
-
 }
+location_types = ['Indoor', 'Outdoor']
+Models = ['Resnet_enc','Unet','Densenet_enc']
+for location in location_types:
+    for model in Models:
+        ckpt_dir = f"{CWD}/ckpt/{CKPT_FILE_NAMES[location][model]}"
+        load_model(CKPT_FILE_NAMES[location][model], MODEL_CLASSES[location][model])
 
-def load_model(ckpt, model, optimizer=None):
-    ckpt_dict = torch.load(ckpt, map_location='cpu')
-    # keep backward compatibility
-    if 'model' not in ckpt_dict and 'optimizer' not in ckpt_dict:
-        state_dict = ckpt_dict
-    else:
-        state_dict = ckpt_dict['model']
-    weights = {}
-    for key, value in state_dict.items():
-        if key.startswith('module.'):
-            weights[key[len('module.'):]] = value
-        else:
-            weights[key] = value
-
-    model.load_state_dict(weights)
-
-    if optimizer is not None:
-        optimizer_state = ckpt_dict['optimizer']
-        optimizer.load_state_dict(optimizer_state)
 
 
 def predict(location, model_name, img):
-    ckpt_dir = f"{CWD}/ckpt/{CKPT_FILE_NAMES[location][model_name]}"
-    if location == 'nyu':
-        max_depth = 10
-    else:
-        max_depth = 80
-    model = MODEL_CLASSES[location][model_name](max_depth).to(DEVICE)
-    load_model(ckpt_dir,model)
+    # ckpt_dir = f"{CWD}/ckpt/{CKPT_FILE_NAMES[location][model_name]}"
+    # if location == 'nyu':
+    #     max_depth = 10
+    # else:
+    #     max_depth = 80
+    # model = MODEL_CLASSES[location][model_name](max_depth).to(DEVICE)
+    model = MODEL_CLASSES[location][model_name].to(DEVICE)
+    # load_model(ckpt_dir,model)
     # print(img.shape)
     # assert False 
     if img.shape ==  (375,1242,3):