add possibility to use two networks

app.py

@@ -13,8 +13,8 @@ from models.preprocessing import *
 from models.misc import nested_tensor_from_tensor_list
 
 
-model = create_letr()
-
+model = create_letr('resnet50/checkpoint0024.pth')
+model101 = create_letr('resnet101/checkpoint0024.pth')
 # PREPARE PREPROCESSING
 # transform_test = transforms.Compose([
 #     transforms.Resize((test_size)),
@@ -38,7 +38,7 @@ normalize_1100 = Compose([
 ])
 
 
-def predict(inp, size):
+def predict(inp, size, model_name):
     image = Image.fromarray(inp.astype('uint8'), 'RGB')
     h, w = image.height, image.width
     orig_size = torch.as_tensor([int(h), int(w)])
@@ -52,7 +52,10 @@ def predict(inp, size):
     inputs = nested_tensor_from_tensor_list([img])
 
     with torch.no_grad():
-        outputs = model(inputs)[0]
+        if model_name == 'resnet101':
+            outputs = model101(inputs)[0]
+        else:
+            outputs = model(inputs)[0]
 
     draw_fig(image, outputs, orig_size)
 
@@ -62,6 +65,7 @@ def predict(inp, size):
 inputs = [
     gr.inputs.Image(),
     gr.inputs.Radio(["256", "512", "1100"]),
+    gr.inputs.Radio(["resnet50", "resnet101"]),
 ]
 outputs = gr.outputs.Image()
 gr.Interface(
@@ -69,8 +73,8 @@ gr.Interface(
     inputs=inputs,
     outputs=outputs,
     examples=[
-        ["demo.png", '256'], 
-        ["tappeto-per-calibrazione.jpg", '256']
+        ["demo.png", '256', "resnet50"], 
+        ["tappeto-per-calibrazione.jpg", '256', "resnet50"]
     ],
     title="LETR",
     description="Model for line detection..."

resnet101/checkpoint0024.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab07a2ddf5e088540941c755e2cccc66e019dcf94b3ef488bab25f5a76490bb9
+size 457215616

test.py

@@ -7,9 +7,9 @@ from models.letr import build
 from models.misc import nested_tensor_from_tensor_list
 from models.preprocessing import Compose, ToTensor, Resize, Normalize
 
-def create_letr():
+def create_letr(path):
     # obtain checkpoints
-    checkpoint = torch.load('checkpoint0024.pth', map_location='cpu')
+    checkpoint = torch.load(path, map_location='cpu')
 
     # load model
     args = checkpoint['args']
@@ -44,7 +44,7 @@ def draw_fig(image, outputs, orig_size):
             draw.line((x1, y1, x2, y2), fill=500)
 
 if __name__ == '__main__':
-    model = create_letr()
+    model = create_letr('resnet50/checkpoint0024.pth')
 
     test_size = 256
     normalize = Compose([