Initial commit.

app.py

@@ -149,7 +149,7 @@ def get_transformations(input_img):
     return transforms.Compose(trans_list)
 
 
-def predict(input_img, cs_mapping):
+def predict(input_img):
     input_img_pil = Image.open(input_img)
     transform = get_transformations(input_img_pil)
     input_img = transform(input_img_pil)
@@ -168,15 +168,13 @@ def predict(input_img, cs_mapping):
     drawing_cs = map2cs(segmentation_remap)
     print('CS colors done.')
 
-    if cs_mapping:
-        drawing = drawing_cs
-    else:
-        drawing = drawing_pseudo
+    drawing_cs = transforms.ToPILImage()(drawing_cs).resize(input_img_pil.size)
+    drawing_blend_cs = blend_images(input_img_pil, drawing_cs)
 
-    drawing = transforms.ToPILImage()(drawing).resize(input_img_pil.size)
-    drawing_blend = blend_images(input_img_pil, drawing)
+    drawing_pseudo = transforms.ToPILImage()(drawing_pseudo).resize(input_img_pil.size)
+    drawing_blend_pseudo = blend_images(input_img_pil, drawing_pseudo)
 
-    return drawing_blend
+    return drawing_blend_pseudo, drawing_blend_cs
 
 
 title = "Drive&Segment"