Initial commit.

app.py

@@ -4,12 +4,11 @@ import requests
 import torch
 import yaml
 from PIL import Image
-from torchvision import transforms
-
 from segmenter_model import utils
 from segmenter_model.factory import create_segmenter
 from segmenter_model.fpn_picie import PanopticFPN
 from segmenter_model.utils import colorize_one, map2cs
+from torchvision import transforms
 
 # WEIGHTS = './weights/segmenter.pth
 WEIGHTS = './weights/segmenter_nusc.pth'
@@ -155,18 +154,12 @@ def predict(input_img):
     input_img = transform(input_img_pil)
     input_img = torch.unsqueeze(input_img, 0)
 
-    print('Loaded and prepaded image.')
-
     with torch.no_grad():
         segmentation = segment_segmenter(input_img, model, window_size, window_stride).squeeze().detach()
-        print('Segmented image.')
         segmentation_remap = remap(segmentation)
-        print('Remapped image.')
 
     drawing_pseudo = colorize_one(segmentation_remap)
-    print('Pseudo colors done.')
     drawing_cs = map2cs(segmentation_remap)
-    print('CS colors done.')
 
     drawing_cs = transforms.ToPILImage()(drawing_cs).resize(input_img_pil.size)
     drawing_blend_cs = blend_images(input_img_pil, drawing_cs)
@@ -181,6 +174,7 @@ title = "Drive&Segment"
 description = 'Gradio Demo accompanying paper "Drive&Segment: Unsupervised Semantic Segmentation of Urban Scenes via Cross-modal Distillation"\nBecause of the CPU-only inference, it might take up to 20s for large images.\nRight now, I use the Segmenter model trained on nuScenes and with 256x256 patches (for the sake of speed).'
 # article = "<p style='text-align: center'><a href='TODO' target='_blank'>Project Page</a> | <a href='codelink' target='_blank'>Github</a></p>"
 examples = ['examples/img5.jpeg', 'examples/100.jpeg', 'examples/39076.jpeg', 'examples/img1.jpg', 'examples/snow1.jpg']
+examples += ['examples/cs{}.png'.format(i) for i in range(1, 6)]
 
 # predict(examples[0])