add radio

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+*.pyc
+venv

app.py

@@ -16,7 +16,6 @@ from models.misc import nested_tensor_from_tensor_list
 model = create_letr()
 
 # PREPARE PREPROCESSING
-test_size = 1100
 # transform_test = transforms.Compose([
 #     transforms.Resize((test_size)),
 #     transforms.ToTensor(),
@@ -25,16 +24,31 @@ test_size = 1100
 normalize = Compose([
         ToTensor(),
         Normalize([0.538, 0.494, 0.453], [0.257, 0.263, 0.273]),
-        Resize([test_size]),
+        Resize([256]),
+])
+normalize_512 = Compose([
+        ToTensor(),
+        Normalize([0.538, 0.494, 0.453], [0.257, 0.263, 0.273]),
+        Resize([512]),
+])
+normalize_1100 = Compose([
+        ToTensor(),
+        Normalize([0.538, 0.494, 0.453], [0.257, 0.263, 0.273]),
+        Resize([1100]),
 ])
 
 
-def predict(inp):
+def predict(inp, size):
     image = Image.fromarray(inp.astype('uint8'), 'RGB')
     h, w = image.height, image.width
     orig_size = torch.as_tensor([int(h), int(w)])
 
-    img = normalize(image)
+    if size == '1100':
+        img = normalize_1100(image)
+    elif size == '512':
+        img = normalize_512(image)
+    else:
+        img = normalize(image)
     inputs = nested_tensor_from_tensor_list([img])
 
     with torch.no_grad():
@@ -45,13 +59,19 @@ def predict(inp):
     return image
 
 
-inputs = gr.inputs.Image()
+inputs = [
+    gr.inputs.Image(),
+    gr.inputs.Radio(["256", "512", "1100"]),
+]
 outputs = gr.outputs.Image()
 gr.Interface(
     fn=predict,
     inputs=inputs,
     outputs=outputs,
-    examples=["demo.png", "tappeto-per-calibrazione.jpg"],
+    examples=[
+        ["demo.png", '256'], 
+        ["tappeto-per-calibrazione.jpg", '256']
+    ],
     title="LETR",
     description="Model for line detection..."
 ).launch()