Add labelme_json_to_dataset utility script

labelme/__init__.py

@@ -0,0 +1 @@
+from labelme import utils

labelme/utils.py

@@ -0,0 +1,94 @@
+import base64
+import cStringIO as StringIO
+
+import matplotlib.pyplot as plt
+import numpy as np
+import PIL.Image
+import PIL.ImageDraw
+import skimage.color
+
+
+def labelcolormap(N=256):
+
+    def bitget(byteval, idx):
+        return ((byteval & (1 << idx)) != 0)
+
+    cmap = np.zeros((N, 3))
+    for i in xrange(0, N):
+        id = i
+        r, g, b = 0, 0, 0
+        for j in xrange(0, 8):
+            r = np.bitwise_or(r, (bitget(id, 0) << 7-j))
+            g = np.bitwise_or(g, (bitget(id, 1) << 7-j))
+            b = np.bitwise_or(b, (bitget(id, 2) << 7-j))
+            id = (id >> 3)
+        cmap[i, 0] = r
+        cmap[i, 1] = g
+        cmap[i, 2] = b
+    cmap = cmap.astype(np.float32) / 255
+    return cmap
+
+
+def img_b64_to_array(img_b64):
+    f = StringIO.StringIO()
+    f.write(base64.b64decode(img_b64))
+    img_arr = np.array(PIL.Image.open(f))
+    return img_arr
+
+
+def polygons_to_mask(img_shape, polygons):
+    mask = np.zeros(img_shape[:2], dtype=np.uint8)
+    mask = PIL.Image.fromarray(mask)
+    xy = map(tuple, polygons)
+    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
+    mask = np.array(mask, dtype=bool)
+    return mask
+
+
+def draw_label(label, img, label_names):
+    plt.subplots_adjust(left=0, right=1, top=1, bottom=0,
+                        wspace=0, hspace=0)
+    plt.margins(0, 0)
+    plt.gca().xaxis.set_major_locator(plt.NullLocator())
+    plt.gca().yaxis.set_major_locator(plt.NullLocator())
+
+    cmap = labelcolormap(len(label_names))
+    label_viz = skimage.color.label2rgb(label, img, bg_label=0)
+    plt.imshow(label_viz)
+
+    plt_handlers = []
+    plt_titles = []
+    for label_value, label_name in enumerate(label_names):
+        fc = cmap[label_value]
+        p = plt.Rectangle((0, 0), 1, 1, fc=fc)
+        plt_handlers.append(p)
+        plt_titles.append(label_name)
+    plt.legend(plt_handlers, plt_titles, loc='lower right', framealpha=.5)
+
+    f = StringIO.StringIO()
+    plt.savefig(f, bbox_inches='tight', pad_inches=0)
+    plt.cla()
+
+    img = np.array(PIL.Image.open(f))[:, :, :3]
+    return img
+
+
+def labelme_shapes_to_label(img_shape, shapes):
+    label_name_to_val = {'background': 0}
+    lbl = np.zeros(img_shape[:2], dtype=np.int32)
+    for shape in sorted(shapes, key=lambda x: x['label']):
+        polygons = shape['points']
+        label_name = shape['label']
+        if label_name in label_name_to_val:
+            label_value = label_name_to_val[label_name]
+        else:
+            label_value = len(label_name_to_val)
+            label_name_to_val[label_name] = label_value
+        mask = polygons_to_mask(img_shape[:2], polygons)
+        lbl[mask] = label_value
+
+    lbl_names = [None] * (max(label_name_to_val.values()) + 1)
+    for label_name, label_value in label_name_to_val.items():
+        lbl_names[label_value] = label_name
+
+    return lbl, lbl_names

labelme_r

@@ -1,7 +0,0 @@
-#!/usr/bin/env python
-
-import labelme.app
-
-
-if __name__ == '__main__':
-    labelme.app.main()

scripts/labelme_draw_json

@@ -2,34 +2,9 @@
 
 import argparse
 import json
-import base64
-from cStringIO import StringIO
-import PIL.Image
-import PIL.ImageDraw
 import matplotlib.pyplot as plt
-import numpy as np
-from skimage.color import label2rgb
 
-
-def labelcolormap(N=256):
-
-    def bitget(byteval, idx):
-        return ((byteval & (1 << idx)) != 0)
-
-    cmap = np.zeros((N, 3))
-    for i in xrange(0, N):
-        id = i
-        r, g, b = 0, 0, 0
-        for j in xrange(0, 8):
-            r = np.bitwise_or(r, (bitget(id, 0) << 7-j))
-            g = np.bitwise_or(g, (bitget(id, 1) << 7-j))
-            b = np.bitwise_or(b, (bitget(id, 2) << 7-j))
-            id = (id >> 3)
-        cmap[i, 0] = r
-        cmap[i, 1] = g
-        cmap[i, 2] = b
-    cmap = cmap.astype(np.float32) / 255
-    return cmap
+from labelme import utils
 
 
 def main():
@@ -41,55 +16,12 @@ def main():
 
     data = json.load(open(json_file))
 
-    # string -> numpy.ndarray
-    f = StringIO()
-    f.write(base64.b64decode(data['imageData']))
-    img = np.array(PIL.Image.open(f))
-
-    target_names = {'background': 0}
-    label = np.zeros(img.shape[:2], dtype=np.int32)
-    for shape in data['shapes']:
-        # get label value
-        label_value = target_names.get(shape['label'])
-        if label_value is None:
-            label_value = len(target_names)
-            target_names[shape['label']] = label_value
-        # polygon -> mask
-        mask = np.zeros(img.shape[:2], dtype=np.uint8)
-        mask = PIL.Image.fromarray(mask)
-        xy = map(tuple, shape['points'])
-        PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
-        mask = np.array(mask)
-        # fill label value
-        label[mask == 1] = label_value
-
-    cmap = labelcolormap(len(target_names))
-    label_viz = label2rgb(label, img, colors=cmap)
-    label_viz[label == target_names['background']] = 0
-
-    plt.subplots_adjust(left=0, right=1, top=1, bottom=0,
-                        wspace=0, hspace=0)
-    plt.margins(0, 0)
-    plt.gca().xaxis.set_major_locator(plt.NullLocator())
-    plt.gca().yaxis.set_major_locator(plt.NullLocator())
-
-    plt.subplot(121)
-    plt.axis('off')
-    plt.imshow(img)
+    img = utils.img_b64_to_array(data['imageData'])
+    lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes'])
 
-    plt.subplot(122)
-    plt.axis('off')
-    plt.imshow(label_viz)
-    # plot target_names legend
-    plt_handlers = []
-    plt_titles = []
-    for l_name, l_value in target_names.items():
-        fc = cmap[l_value]
-        p = plt.Rectangle((0, 0), 1, 1, fc=fc)
-        plt_handlers.append(p)
-        plt_titles.append(l_name)
-    plt.legend(plt_handlers, plt_titles, loc='lower right', framealpha=.5)
+    lbl_viz = utils.draw_label(lbl, img, lbl_names)
 
+    plt.imshow(lbl_viz)
     plt.show()
 
 

scripts/labelme_json_to_dataset

@@ -0,0 +1,45 @@
+#!/usr/bin/python
+
+import argparse
+import json
+import os
+import os.path as osp
+
+import PIL.Image
+import yaml
+
+from labelme import utils
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('json_file')
+    args = parser.parse_args()
+
+    json_file = args.json_file
+
+    out_dir = osp.basename(json_file).replace('.', '_')
+    out_dir = osp.join(osp.dirname(json_file), out_dir)
+    os.mkdir(out_dir)
+
+    data = json.load(open(json_file))
+
+    img = utils.img_b64_to_array(data['imageData'])
+    lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes'])
+
+    lbl_viz = utils.draw_label(lbl, img, lbl_names)
+
+    PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
+    PIL.Image.fromarray(lbl).save(osp.join(out_dir, 'label.png'))
+    PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
+
+    info = dict(label_names=lbl_names)
+
+    with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
+        yaml.safe_dump(info, f, default_flow_style=False)
+
+    print('wrote data to %s' % out_dir)
+
+
+if __name__ == '__main__':
+    main()

setup.py

@@ -67,5 +67,5 @@ setup(
     ],
     package_data={'labelme': ['icons/*', 'resources.qrc']},
     entry_points={'console_scripts': ['labelme=labelme.app:main']},
-    scripts=['scripts/labelme_draw_json'],
+    scripts=['scripts/labelme_draw_json', 'scripts/labelme_json_to_dataset'],
 )