add models

models/DocScanner/inference.py

@@ -37,12 +37,12 @@ class Net(nn.Module):
         return bm, msk
 
 
-def reload_seg_model(model, path=""):
+def reload_seg_model(cuda, model, path=""):
     if not bool(path):
         return model
     else:
         model_dict = model.state_dict()
-        pretrained_dict = torch.load(path, map_location="cuda:0")
+        pretrained_dict = torch.load(path, map_location=cuda)
         pretrained_dict = {
             k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict
         }
@@ -52,12 +52,12 @@ def reload_seg_model(model, path=""):
         return model
 
 
-def reload_rec_model(model, path=""):
+def reload_rec_model(cuda, model, path=""):
     if not bool(path):
         return model
     else:
         model_dict = model.state_dict()
-        pretrained_dict = torch.load(path, map_location="cuda:0")
+        pretrained_dict = torch.load(path, map_location=cuda)
         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
         model_dict.update(pretrained_dict)
         model.load_state_dict(model_dict)

rect_main.py

@@ -33,8 +33,8 @@ def load_geotrp_model(cuda, path=""):
 def load_docscanner_model(cuda, path_l="", path_m=""):
 
     net = DocScanner.Net().to(cuda)
-    DocScanner.reload_seg_model(net.msk, path_m)
-    DocScanner.reload_rec_model(net.bm, path_l)
+    DocScanner.reload_seg_model(cuda, net.msk, path_m)
+    DocScanner.reload_rec_model(cuda, net.bm, path_l)
     net.eval()
 
     return net