Update app.py

app.py

@@ -20,6 +20,35 @@ import gradio as gr
 
 example_img_list = ['51_1 copy.png','48_2 copy.png','25.jpg'] 
 
+def reload_model(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cpu')
+        # print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
+        # print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+
+
+def reload_segmodel(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cpu')
+        # print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}
+        # print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+
 class GeoTr_Seg(nn.Module):
     def __init__(self):
         super(GeoTr_Seg, self).__init__()
@@ -36,6 +65,9 @@ class GeoTr_Seg(nn.Module):
         
         return bm
 
+
+        
+
 # Initialize models
 GeoTr_Seg_model = GeoTr_Seg()
 #IllTr_model = IllTr()