bundle demo image + label and save output

shelhamer · Jan 23, 2018 · 31e60a7 · 31e60a7
1 parent 230e02c
commit 31e60a7
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 1 deletion.
diff --git a/demo/image.jpg b/demo/image.jpg
diff --git a/demo/label.png b/demo/label.png
diff --git a/demo/output.png b/demo/output.png
diff --git a/demo/visualization.jpg b/demo/visualization.jpg
diff --git a/infer.py b/infer.py
@@ -2,9 +2,12 @@
 from PIL import Image
 
 import caffe
+import vis
+
+# the demo image is "2007_000129" from PASCAL VOC
 
 # load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
-im = Image.open('pascal/VOC2010/JPEGImages/2007_000129.jpg')
+im = Image.open('demo/image.jpg')
 in_ = np.array(im, dtype=np.float32)
 in_ = in_[:,:,::-1]
 in_ -= np.array((104.00698793,116.66876762,122.67891434))
@@ -18,3 +21,10 @@
 # run net and take argmax for prediction
 net.forward()
 out = net.blobs['score'].data[0].argmax(axis=0)
+
+# visualize segmentation in PASCAL VOC colors
+voc_palette = vis.make_palette(21)
+out_im = Image.fromarray(vis.color_seg(out, voc_palette))
+out_im.save('demo/output.png')
+masked_im = Image.fromarray(vis.vis_seg(im, out, voc_palette))
+masked_im.save('demo/visualization.jpg')