Merge pull request #224 from jina-ai/feat-crafter-image-200

add image crafters

jina/drivers/craft.py

@@ -17,10 +17,31 @@ def __call__(self, *args, **kwargs):
             if not d.chunks:
                 no_chunk_docs.append(d.doc_id)
                 continue
+            _chunks_to_add = []
             for c in d.chunks:
                 ret = self.exec_fn(**pb_obj2dict(c, self.exec.required_keys))
-                for k, v in ret.items():
-                    setattr(c, k, v)
+                if isinstance(ret, dict):
+                    for k, v in ret.items():
+                        setattr(c, k, v)
+                    continue
+                if isinstance(ret, list):
+                    for chunk_dict in ret:
+                        _chunks_to_add.append(chunk_dict)
+            if len(_chunks_to_add) > 0:
+                for c_dict in _chunks_to_add:
+                    c = d.chunks.add()
+                    for k, v in c_dict.items():
+                        if k == 'blob':
+                            c.blob.CopyFrom(array2blob(v))
+                        elif k == 'chunk_id':
+                            self.logger.warning(f'you are assigning a chunk_id in in {self.exec.__class__}, '
+                                                f'is it intentional? chunk_id will be override by {self.__class__} '
+                                                f'anyway')
+                        else:
+                            setattr(c, k, v)
+                    c.length = len(_chunks_to_add) + len(d.chunks)
+                    c.chunk_id = random.randint(0, ctypes.c_uint(-1).value)
+            d.length = len(_chunks_to_add) + len(d.chunks)
 
         if no_chunk_docs:
             self.logger.warning('these docs contain no chunk: %s' % no_chunk_docs)
@@ -46,7 +67,7 @@ class SegmentDriver(BaseExecutableDriver):
         no need to self-assign it in your segmenter
     """
 
-    def __init__(self, first_chunk_id: int = 0, random_chunk_id: bool = False, *args, **kwargs):
+    def __init__(self, first_chunk_id: int = 0, random_chunk_id: bool = True, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.first_chunk_id = first_chunk_id
         self.random_chunk_id = random_chunk_id
@@ -62,7 +83,8 @@ def __call__(self, *args, **kwargs):
                             c.blob.CopyFrom(array2blob(v))
                         elif k == 'chunk_id':
                             self.logger.warning(f'you are assigning a chunk_id in in {self.exec.__class__}, '
-                                                f'is it intentional? chunk_id will be override by {self.__class__} anyway')
+                                                f'is it intentional? chunk_id will be override by {self.__class__} '
+                                                f'anyway')
                         else:
                             setattr(c, k, v)
                     c.length = len(ret)

jina/executors/crafters/image/__init__.py

@@ -0,0 +1,110 @@
+import numpy as np
+import numbers
+
+from .. import BaseChunkCrafter
+from typing import Tuple, Union, Iterable
+
+
+class ImageChunkCrafter(BaseChunkCrafter):
+    """
+    :class:`ImageChunkCrafter` provides the basic functions for processing image data on chunk-level.
+
+    .. warning::
+        :class:'ImageChunkCrafter' is intended to be used internally.
+
+    """
+    def __init__(self, channel_axis: int = -1, *args, **kwargs):
+        """
+
+        :param channel_axis: the axis id of the color channel, ``-1`` indicates the color channel info at the last axis
+        """
+        super().__init__(*args, **kwargs)
+        self.channel_axis = channel_axis
+
+    def check_channel_axis(self, img: 'np.ndarray') -> 'np.ndarray':
+        """
+        Ensure the color channel axis is the last axis.
+        """
+        if self.channel_axis == -1:
+            return img
+        return np.moveaxis(img, self.channel_axis, -1)
+
+    def _load_image(self, blob: 'np.ndarray'):
+        from PIL import Image
+        img = self.check_channel_axis(blob)
+        return Image.fromarray(img.astype('uint8'))
+
+    @staticmethod
+    def _resize_short(img, target_size: Union[Tuple[int], int], how: str = 'LANCZOS'):
+        """
+        Resize the input :py:mod:`PIL` image.
+
+        :param img: :py:mod:`PIL.Image`, the image to be resized
+        :param target_size: desired output size. If size is a sequence like (h, w), the output size will be matched to
+            this. If size is an int, the smaller edge of the image will be matched to this number maintain the aspect
+            ratio.
+        :param how: the interpolation method. Valid values include `NEAREST`, `BILINEAR`, `BICUBIC`, and `LANCZOS`.
+            Default is `LANCZOS`. Please refer to `PIL.Image` for detaisl.
+        """
+        import PIL.Image as Image
+        assert isinstance(img, Image.Image), 'img must be a PIL.Image'
+        if isinstance(target_size, int):
+            percent = float(target_size) / min(img.size[0], img.size[1])
+            target_w = int(round(img.size[0] * percent))
+            target_h = int(round(img.size[1] * percent))
+        elif isinstance(target_size, Tuple) and len(target_size) == 2:
+            target_h, target_w = target_size
+        else:
+            raise ValueError('target_size should be an integer or a tuple of two integers: {}'.format(target_size))
+        img = img.resize((target_w, target_h), getattr(Image, how))
+        return img
+
+    @staticmethod
+    def _crop_image(img, target_size: Union[Tuple[int], int], top: int = None, left: int = None, how: str = 'precise'):
+        """
+        Crop the input :py:mod:`PIL` image.
+
+        :param img: :py:mod:`PIL.Image`, the image to be resized
+        :param target_size: desired output size. If size is a sequence like
+            (h, w), the output size will be matched to this. If size is an int,
+            the output will have the same height and width as the `target_size`.
+        :param top: the vertical coordinate of the top left corner of the crop box.
+        :param left: the horizontal coordinate of the top left corner of the crop box.
+        :param how: the way of cropping. Valid values include `center`, `random`, and, `precise`. Default is `precise`.
+            - `center`: crop the center part of the image
+            - `random`: crop a random part of the image
+            - `precise`: crop the part of the image specified by the crop box with the given ``top`` and ``left``.
+            .. warning:: When `precise` is used, ``top`` and ``left`` must be fed valid value.
+
+        """
+        import PIL.Image as Image
+        assert isinstance(img, Image.Image), 'img must be a PIL.Image'
+        img_w, img_h = img.size
+        if isinstance(target_size, int):
+            target_h = target_w = target_size
+        elif isinstance(target_size, Tuple) and len(target_size) == 2:
+            target_h, target_w = target_size
+        else:
+            raise ValueError('target_size should be an integer or a tuple of two integers: {}'.format(target_size))
+        w_beg = left
+        h_beg = top
+        if how == 'center':
+            w_beg = int((img_w - target_w) / 2)
+            h_beg = int((img_h - target_h) / 2)
+        elif how == 'random':
+            w_beg = np.random.randint(0, img_w - target_w + 1)
+            h_beg = np.random.randint(0, img_h - target_h + 1)
+        elif how == 'precise':
+            assert (w_beg is not None and h_beg is not None)
+            assert (0 <= w_beg <= (img_w - target_w)), 'left must be within [0, {}]: {}'.format(img_w - target_w, w_beg)
+            assert (0 <= h_beg <= (img_h - target_h)), 'top must be within [0, {}]: {}'.format(img_h - target_h, h_beg)
+        else:
+            raise ValueError('unknown input how: {}'.format(how))
+        if not isinstance(w_beg, int):
+            raise ValueError('left must be int number between 0 and {}: {}'.format(img_w, left))
+        if not isinstance(h_beg, int):
+            raise ValueError('top must be int number between 0 and {}: {}'.format(img_h, top))
+        w_end = w_beg + target_w
+        h_end = h_beg + target_h
+        img = img.crop((w_beg, h_beg, w_end, h_end))
+        return img