Show class_id in Visualization

huang-salute · Mar 31, 2018 · 65acf30 · 65acf30
1 parent e70bcae
commit 65acf30
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/semantic_segmentation/labelme2voc.py b/examples/semantic_segmentation/labelme2voc.py
@@ -89,8 +89,10 @@ def main():
             # lbl_pil.putpalette((colormap * 255).flatten())
             lbl_pil.save(out_lbl_file)
 
+            label_names = ['%d: %s' % (cls_id, cls_name)
+                           for cls_id, cls_name in enumerate(class_names)]
             viz = labelme.utils.draw_label(
-                lbl, img, class_names, colormap=colormap)
+                lbl, img, label_names, colormap=colormap)
             PIL.Image.fromarray(viz).save(out_viz_file)