Class ToMode

imperative/python/megengine/data/transform/vision/transform.py:149–175

r"""Change input data to a target mode. For example, most transforms use HWC mode image, while the neural network might use CHW mode input tensor. Args: mode: output mode of input. Default: "CHW" order: the same with :class:`VisionTransform`

class ToMode(VisionTransform):
    r"""Change input data to a target mode.
    For example, most transforms use HWC mode image,
    while the neural network might use CHW mode input tensor.

    Args:
        mode: output mode of input. Default: "CHW"
        order: the same with :class:`VisionTransform`
    """

    def __init__(self, mode="CHW", *, order=None):
        super().__init__(order)
        assert mode in ["CHW"], "unsupported mode: {}".format(mode)
        self.mode = mode

    def _apply_image(self, image):
        if self.mode == "CHW":
            return np.ascontiguousarray(np.rollaxis(image, 2))
        return image

    def _apply_coords(self, coords):
        return coords

    def _apply_mask(self, mask):
        if self.mode == "CHW":
            return np.ascontiguousarray(np.rollaxis(mask, 2))
        return mask


class Compose(VisionTransform):