Pārlūkot izejas kodu

Add img_b64_to_arr, img_arr_to_b64

Kentaro Wada 7 gadi atpakaļ
vecāks
revīzija
24259a5810
5 mainītis faili ar 41 papildinājumiem un 10 dzēšanām
  1. 10 6
      .gitignore
  2. 15 0
      labelme/utils.py
  3. 1 0
      tests/data/apc2016_obj3.jpg
  4. 1 0
      tests/data/apc2016_obj3.json
  5. 14 4
      tests/test_utils.py

+ 10 - 6
.gitignore

@@ -1,8 +1,12 @@
-.cache/
-build/
-dist/
+/.cache/
+/.pytest_cache/
+
+/build/
+/dist/
+/*.egg-info/
+
 *.py[cdo]
-*.egg-info/
 
-icons/.DS_Store
-labelme/resources.py
+.DS_Store
+
+/labelme/resources.py

+ 15 - 0
labelme/utils.py

@@ -57,12 +57,27 @@ def label2rgb(lbl, img=None, n_labels=None, alpha=0.3, thresh_suppress=0):
 
 
 def img_b64_to_array(img_b64):
+    warnings.warn('img_ba64_to_array is deprecated. '
+                  'Please use img_b64_to_arr.')
+    return img_b64_to_arr(img_b64)
+
+
+def img_b64_to_arr(img_b64):
     f = io.BytesIO()
     f.write(base64.b64decode(img_b64))
     img_arr = np.array(PIL.Image.open(f))
     return img_arr
 
 
+def img_arr_to_b64(img_arr):
+    img_pil = PIL.Image.fromarray(img_arr)
+    f = io.BytesIO()
+    img_pil.save(f, format='PNG')
+    img_bin = f.getvalue()
+    img_b64 = base64.encodestring(img_bin)
+    return img_b64
+
+
 def polygons_to_mask(img_shape, polygons):
     mask = np.zeros(img_shape[:2], dtype=np.uint8)
     mask = PIL.Image.fromarray(mask)

+ 1 - 0
tests/data/apc2016_obj3.jpg

@@ -0,0 +1 @@
+../../examples/single_image/apc2016_obj3.jpg

+ 1 - 0
tests/data/apc2016_obj3.json

@@ -0,0 +1 @@
+../../examples/single_image/apc2016_obj3.json

+ 14 - 4
tests/test_utils.py

@@ -2,17 +2,27 @@ import json
 import os.path as osp
 
 import numpy as np
+import PIL.Image
 
-from labelme import utils
+import labelme
 
 
 here = osp.dirname(osp.abspath(__file__))
+data_dir = osp.join(here, 'data')
 
 
-def test_img_b64_to_array():
-    json_file = osp.join(here, '../examples/single_image/apc2016_obj3.json')
+def test_img_b64_to_arr():
+    json_file = osp.join(data_dir, 'apc2016_obj3.json')
     data = json.load(open(json_file))
     img_b64 = data['imageData']
-    img = utils.img_b64_to_array(img_b64)
+    img = labelme.utils.img_b64_to_arr(img_b64)
     assert img.dtype == np.uint8
     assert img.shape == (907, 1210, 3)
+
+
+def test_img_arr_to_b64():
+    img_file = osp.join(data_dir, 'apc2016_obj3.jpg')
+    img_arr = np.asarray(PIL.Image.open(img_file))
+    img_b64 = labelme.utils.img_arr_to_b64(img_arr)
+    img_arr2 = labelme.utils.img_b64_to_arr(img_b64)
+    np.testing.assert_allclose(img_arr, img_arr2)