Initial commit.

app.py

@@ -14,9 +14,10 @@ from segmenter_model.utils import colorize_one, map2cs
 # WEIGHTS = './weights/segmenter.pth
 WEIGHTS = './weights/segmenter_nusc.pth'
 FULL = True
+ALPHA = 0.5
 
 
-def blend_images(bg, fg, alpha=0.3):
+def blend_images(bg, fg, alpha=ALPHA):
     fg = fg.convert('RGBA')
     bg = bg.convert('RGBA')
     blended = Image.blend(bg, fg, alpha=alpha)
@@ -135,10 +136,12 @@ download_weights()
 model, window_size, window_stride, im_size = create_model()
 
 
-def get_transformations():
+def get_transformations(input_img):
     trans_list = [transforms.ToTensor()]
 
-    if im_size != 1024:
+    shorter_input_size = min(input_img.size)
+
+    if im_size != 1024 or shorter_input_size < im_size:
         trans_list.append(transforms.Resize(im_size))
 
     trans_list.append(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
@@ -148,7 +151,7 @@ def get_transformations():
 
 def predict(input_img, cs_mapping):
     input_img_pil = Image.open(input_img)
-    transform = get_transformations()
+    transform = get_transformations(input_img)
     input_img = transform(input_img_pil)
     input_img = torch.unsqueeze(input_img, 0)
 
@@ -186,7 +189,7 @@ examples = [['examples/img5.jpeg', True], ['examples/100.jpeg', True], ['example
 
 iface = gr.Interface(predict, [gr.inputs.Image(type='filepath'), gr.inputs.Checkbox(label="Cityscapes mapping")],
                      "image", title=title, description=description,
-                     examples=examples)
+                     examples=examples, allow_screenshot=True)
 # iface = gr.Interface(predict, gr.inputs.Image(type='filepath'),
 #                      "image", title=title, description=description,
 #                      examples=examples)