def _prepare_image()

in python/mxboard/utils.py [0:0]

• 24 lines of code
• 10 McCabe index (conditional complexity)

def _prepare_image(img, nrow=8, padding=2, square_image=False):
    """Given an image of format HW, CHW, or NCHW, returns a image of format HWC.
    If the input is a batch of images, a grid of images is made by stitching them together.
    If data type is float, values must be in the range [0, 1], and then they are rescaled to
    range [0, 255]. If data type is 'uint8`, values are unchanged.
    """
    if isinstance(img, np.ndarray):
        img = nd.array(img, dtype=img.dtype, ctx=current_context())
    if not isinstance(img, NDArray):
        raise TypeError('expected MXNet NDArray or numpy.ndarray, '
                        'while received type {}'.format(str(type(img))))
    assert img.ndim == 2 or img.ndim == 3 or img.ndim == 4

    if img.dtype == np.uint8:
        return make_image_grid(
            img, nrow=nrow, padding=padding, square_image=square_image).transpose((1, 2, 0))
    elif img.dtype == np.float32 or img.dtype == np.float64:
        min_val = img.min().asscalar()
        max_val = img.max().asscalar()
        if min_val < 0.0:
            raise ValueError('expected non-negative min value from img, '
                             'while received {}'.format(min_val))
        if max_val > 1.0:
            raise ValueError('expected max value from img not greater than 1, '
                             'while received {}'.format(max_val))
        img = make_image_grid(img, nrow=nrow, padding=padding, square_image=square_image) * 255.0
        return img.astype(np.uint8).transpose((1, 2, 0))
    else:
        raise ValueError('expected input image dtype is one of uint8, float32, '
                         'and float64, received dtype {}'.format(str(img.dtype)))