code: __copyright__ = "Copyright (c) 2021 Jina AI Limited. All rights reserved."
  __license__ = "Apache-2.0"
  
  from typing import Tuple, Dict, Union, Iterable
  
  import numpy as np
  
  from jina.executors.decorators import single
  from jina.executors.crafters import BaseCrafter
  
  from .helper import _load_image, _move_channel_axis, _crop_image, _resize_short
  
  
  class ImageNormalizer(BaseCrafter):
      """
      Normalize the image.
  
      :class:`ImageNormalizer` works on doc-level,
          it receives values of file names on the
          doc-level and returns image matrix on the chunk-level
  
      :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 maintaining the aspect ratio.
      :param img_mean: The mean of the images in `RGB` channels.
          Set to `[0.485, 0.456, 0.406]` for the models trained
          on `imagenet` with pytorch backbone.
      :param img_std: the std of the images in `RGB` channels.
          Set to `[0.229, 0.224, 0.225]` for the models trained
          on `imagenet` with pytorch backbone.
      :param resize_dim: the size of images' height and width to be resized to.
          The images are resized before cropping to the output size
      :param channel_axis: the axis id of the color channel,
          ``-1`` indicates the color channel info at the last axis
      """
  
      def __init__(self,
                   target_size: Union[Iterable[int], int] = 224,
                   img_mean: Tuple[float] = (0, 0, 0),
                   img_std: Tuple[float] = (1, 1, 1),
                   resize_dim: int = 256,
                   channel_axis: int = -1,
                   *args,
                   **kwargs):
          """Set Constructor."""
          super().__init__(*args, **kwargs)
          if isinstance(target_size, int):
              self.target_size = target_size
          elif isinstance(target_size, Iterable):
              self.target_size = tuple(target_size)
          else:
              raise ValueError(f'target_size {target_size} should be an integer or tuple/list of 2 integers')
          self.resize_dim = resize_dim
          self.img_mean = np.array(img_mean).reshape((1, 1, 3))
          self.img_std = np.array(img_std).reshape((1, 1, 3))
          self.channel_axis = channel_axis
  
      @single
      def craft(self, blob: 'np.ndarray', *args, **kwargs) -> Dict:
          """
          Normalize the image.
  
          :param blob: the ndarray of the image with the color channel at the last axis
          :return: a chunk dict with the normalized image
          """
          raw_img = _load_image(blob, self.channel_axis)
          _img = self._normalize(raw_img)
          img = _move_channel_axis(_img, -1, self.channel_axis)
          return dict(offset=0, blob=img)
  
      def _normalize(self, img):
          img = _resize_short(img, target_size=self.resize_dim)
          img, _, _ = _crop_image(img, target_size=self.target_size, how='center')
          img = np.array(img).astype('float32')/255
          img -= self.img_mean
          img /= self.img_std
          return img
   import io
  
  import numpy as np
  from PIL import Image
  
  from .. import ImageReader
  
  
  def create_test_image(output_fn, size_width=50, size_height=50):
      from PIL import Image
      image = Image.new('RGB', size=(size_width, size_height), color=(155, 0, 0))
      with open(output_fn, "wb") as f:
          image.save(f, 'jpeg')
  
  
  def test_io_uri():
      crafter = ImageReader()
      tmp_fn = crafter.get_file_from_workspace('test.jpeg')
      img_size = 50
      create_test_image(tmp_fn, size_width=img_size, size_height=img_size)
      test_docs = crafter.craft([None, None], np.stack([tmp_fn, tmp_fn]))
      assert len(test_docs) == 2
      for test_doc in test_docs:
          assert test_doc['blob'].shape == (img_size, img_size, 3)
  
  
  def test_io_buffer():
      crafter = ImageReader()
      tmp_fn = crafter.get_file_from_workspace('test.jpeg')
      img_size = 50
      create_test_image(tmp_fn, size_width=img_size, size_height=img_size)
      image_buffer = io.BytesIO()
      img = Image.open(tmp_fn)
      img.save(image_buffer, format='PNG')
      image_buffer.seek(0)
      test_docs = crafter.craft(np.stack([image_buffer.getvalue(), image_buffer.getvalue()]), [None, None])
      assert len(test_docs) == 2
      for test_doc in test_docs:
          assert test_doc['blob'].shape == (img_size, img_size, 3)
          np.testing.assert_almost_equal(test_doc['blob'], np.array(img).astype('float32'))
  
apis: numpy.array
apis: numpy.stack numpy.array