Repository: mars-project/mars

Issue:
Support reading images data
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

Support loading image files to Mars tensor, it is useful for image processing. Users can use `Mars.imread` to read a single image or multiple images to Mars tensor.


Discussion:

Resolving patch:
diff --git a/mars/tensor/images/__init__.py b/mars/tensor/images/__init__.py
new file mode 100644
--- /dev/null
+++ b/mars/tensor/images/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .imread import imread
diff --git a/mars/tensor/images/imread.py b/mars/tensor/images/imread.py
new file mode 100644
--- /dev/null
+++ b/mars/tensor/images/imread.py
@@ -0,0 +1,99 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+
+from ... import opcodes as OperandDef
+from ...serialize import AnyField
+from ...config import options
+from ...filesystem import open_file, glob, file_size
+from ...utils import ceildiv
+from ..operands import TensorOperandMixin, TensorOperand
+
+
+def _read_image(fpath):
+    return np.asarray(Image.open(fpath))
+
+
+class TensorImread(TensorOperand, TensorOperandMixin):
+    _op_type_ = OperandDef.IMREAD
+
+    _filepath = AnyField('filepath')
+
+    def __init__(self, filepath=None, **kwargs):
+        super().__init__(_filepath=filepath, **kwargs)
+
+    @property
+    def filepath(self):
+        return self._filepath
+
+    @classmethod
+    def tile(cls, op):
+        out_shape = op.outputs[0].shape
+        paths = op.filepath if isinstance(op.filepath, (tuple, list)) else glob(op.filepath)
+        chunk_size = op.outputs[0].extra_params.raw_chunk_size
+        n_chunks = ceildiv(len(paths), chunk_size)
+        if len(paths) > 1:
+            chunks = []
+            splits = []
+            for i in range(n_chunks):
+                chunk_op = op.copy().reset_key()
+                chunk_op._filepath = paths[i * chunk_size: (i + 1) * chunk_size]
+                file_nums = len(chunk_op._filepath)
+                shape = (file_nums,) + out_shape[1:]
+                chunk = chunk_op.new_chunk(None, shape=shape, index=(i,) + (0,) * (len(out_shape) - 1))
+                chunks.append(chunk)
+                splits.append(file_nums)
+            nsplits = (tuple(splits),) + tuple((s,) for s in out_shape[1:])
+        else:
+            chunk_op = op.copy().reset_key()
+            chunks = [chunk_op.new_chunk(None, shape=out_shape)]
+            nsplits = tuple((s,) for s in out_shape)
+        new_op = op.copy()
+        return new_op.new_tensors(None, shape=out_shape, chunks=chunks, nsplits=nsplits)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        if isinstance(op.filepath, list):
+            arrays = np.empty(op.outputs[0].shape)
+            for i, path in enumerate(op.filepath):
+                with open_file(path, 'rb') as f:
+                    arrays[i] = _read_image(f)
+            ctx[op.outputs[0].key] = np.array(arrays)
+        else:
+            with open_file(op.filepath, 'rb') as f:
+                ctx[op.outputs[0].key] = np.array(_read_image(f))
+
+    def __call__(self, shape, chunk_size):
+        return self.new_tensor(None, shape, raw_chunk_size=chunk_size)
+
+
+def imread(path, chunk_size=None):
+    paths = path if isinstance(path, (tuple, list)) else glob(path)
+    with open_file(paths[0], 'rb') as f:
+        sample_data = _read_image(f)
+        img_shape = sample_data.shape
+        img_size = file_size(paths[0])
+    if len(paths) > 1:
+        shape = (len(paths), ) + img_shape
+    else:
+        shape = img_shape
+    if chunk_size is None:
+        chunk_size = int(options.chunk_store_limit / img_size)
+    op = TensorImread(filepath=path, chunk_size=chunk_size)
+    return op(shape=shape, chunk_size=chunk_size)


Test patch:
diff --git a/mars/tensor/images/tests/__init__.py b/mars/tensor/images/tests/__init__.py
new file mode 100644
--- /dev/null
+++ b/mars/tensor/images/tests/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
diff --git a/mars/tensor/images/tests/test_images.py b/mars/tensor/images/tests/test_images.py
new file mode 100644
--- /dev/null
+++ b/mars/tensor/images/tests/test_images.py
@@ -0,0 +1,60 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+import tempfile
+
+import numpy as np
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+
+from mars.tests.core import TestBase
+from mars.tensor.images import imread
+
+
+@unittest.skipIf(not Image, 'Pillow not installed')
+class Test(TestBase):
+    def testImread(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            raws = []
+            for i in range(10):
+                array = np.random.randint(0, 256, 2500 * 3, dtype=np.uint8).reshape((50, 50, 3))
+                raws.append(array)
+                im = Image.fromarray(array)
+                im.save(os.path.join(tempdir, 'random_{}.png'.format(i)))
+
+            t = imread(os.path.join(tempdir, 'random_0.png'))
+            self.assertEqual(t.shape, (50, 50, 3))
+
+            tiled = t.tiles()
+            self.assertEqual(len(tiled.chunks), 1)
+            self.assertEqual(tiled.chunks[0].shape, (50, 50, 3))
+
+            t = imread(os.path.join(tempdir, 'random_*.png'), chunk_size=3)
+            self.assertEqual(t.shape, (10, 50, 50, 3))
+
+            tiled = t.tiles()
+            self.assertEqual(len(tiled.chunks), 4)
+            self.assertEqual(tiled.nsplits, ((3, 3, 3, 1), (50,), (50,), (3,)))
+            self.assertEqual(tiled.chunks[0].index, (0, 0, 0, 0))
+            self.assertEqual(tiled.chunks[0].shape, (3, 50, 50, 3))
+            self.assertEqual(tiled.chunks[1].index, (1, 0, 0, 0))
+            self.assertEqual(tiled.chunks[1].shape, (3, 50, 50, 3))
+            self.assertEqual(tiled.chunks[2].index, (2, 0, 0, 0))
+            self.assertEqual(tiled.chunks[2].shape, (3, 50, 50, 3))
+            self.assertEqual(tiled.chunks[3].index, (3, 0, 0, 0))
+            self.assertEqual(tiled.chunks[3].shape, (1, 50, 50, 3))
diff --git a/mars/tensor/images/tests/test_images_execute.py b/mars/tensor/images/tests/test_images_execute.py
new file mode 100644
--- /dev/null
+++ b/mars/tensor/images/tests/test_images_execute.py
@@ -0,0 +1,58 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+import tempfile
+
+import numpy as np
+try:
+    from PIL import Image
+except ImportError:
+    Image = None
+
+from mars.tests.core import TestBase, ExecutorForTest
+from mars.tensor.images import imread
+
+
+@unittest.skipIf(not Image, 'Pillow not installed')
+class Test(TestBase):
+    def setUp(self):
+        super().setUp()
+        self.executor = ExecutorForTest()
+
+    def testImreadExecution(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            raws = []
+            for i in range(10):
+                array = np.random.randint(0, 256, 2500, dtype=np.uint8).reshape((50, 50))
+                raws.append(array)
+                im = Image.fromarray(array)
+                im.save(os.path.join(tempdir, 'random_{}.png'.format(i)))
+            # Single image
+            t = imread(os.path.join(tempdir, 'random_0.png'))
+            res = self.executor.execute_tensor(t, concat=True)[0]
+            np.testing.assert_array_equal(res, raws[0])
+
+            t2 = imread(os.path.join(tempdir, 'random_*.png'))
+            res = self.executor.execute_tensor(t2, concat=True)[0]
+            np.testing.assert_array_equal(np.sort(res, axis=0), np.sort(raws, axis=0))
+
+            t3 = imread(os.path.join(tempdir, 'random_*.png'), chunk_size=4)
+            res = self.executor.execute_tensor(t3, concat=True)[0]
+            np.testing.assert_array_equal(np.sort(res, axis=0), np.sort(raws, axis=0))
+
+            t4 = imread(os.path.join(tempdir, 'random_*.png'), chunk_size=4)
+            res = self.executor.execute_tensor(t4, concat=True)[0]
+            np.testing.assert_array_equal(np.sort(res, axis=0), np.sort(raws, axis=0))