<commit_before>from __future__ import print_function, division, absolute_import

import numpy as np

import imgaug.imgaug as ia
import imgaug.dtypes as iadt


def normalize_images(images):
    if images is None:
        return None
    elif ia.is_np_array(images):
        if images.ndim == 2:
            return images[np.newaxis, ..., np.newaxis]
        elif images.ndim == 3:
            return images[..., np.newaxis]
        else:
            return images
    elif ia.is_iterable(images):
        result = []
        for image in images:
            if image.ndim == 2:
                result.append(image[..., np.newaxis])
            else:
                assert image.ndim == 3
                result.append(image)
        return result
    raise ValueError(
        ("Expected argument 'images' to be any of the following: "
         + "None or array or iterable of array. Got type: %s.") % (
            type(images),)
    )


def normalize_heatmaps(inputs, images=None):
    # TODO get rid of this deferred import
    from imgaug.augmentables.heatmaps import HeatmapsOnImage

    ntype = estimate_heatmaps_norm_type(inputs)
    if ntype == "None":
        return None
    elif ntype == "array[float]":
        assert images is not None
        assert inputs.ndim == 4  # always (N,H,W,C)
        assert len(inputs) == len(images)
        return [HeatmapsOnImage(attr_i, shape=image_i.shape)
                for attr_i, image_i in zip(inputs, images)]
    elif ntype == "HeatmapsOnImage":
        return [inputs]
    elif ntype == "iterable[empty]":
        return None
    elif ntype == "iterable-array[float]":
        assert images is not None
        assert len(inputs) == len(images)
        assert all([attr_i.ndim == 3 for attr_i in inputs])  # all (H,W,C)
        return [HeatmapsOnImage(attr_i, shape=image_i.shape)
                for attr_i, image_i in zip(inputs, images)]
    else:
        assert ntype == "iterable-HeatmapsOnImage"
        return inputs  # len allowed to differ from len of images


def normalize_segmentation_maps(inputs, images=None):
    # TODO get rid of this deferred import
    from imgaug.augmentables.segmaps import SegmentationMapOnImage

    ntype = estimate_segmaps_norm_type(inputs)
    if ntype == "None":
        return None
    elif ntype in ["array[int]", "array[uint]", "array[bool]"]:
        assert images is not None
        assert inputs.ndim == 4  # always (N,H,W,C)
        assert len(inputs) == len(images)
        if ntype == "array[bool]":
            return [SegmentationMapOnImage(attr_i, shape=image_i.shape)
                    for attr_i, image_i in zip(inputs, images)]
        return [SegmentationMapOnImage(
                    attr_i, shape=image_i.shape, nb_classes=1+np.max(attr_i))
                for attr_i, image_i in zip(inputs, images)]
    elif ntype == "SegmentationMapOnImage":
        return [inputs]
    elif ntype == "iterable[empty]":
        return None
    elif ntype in ["iterable-array[int]",
                   "iterable-array[uint]",
                   "iterable-array[bool]"]:
        assert images is not None
        assert len(inputs) == len(images)
        assert all([attr_i.ndim == 3 for attr_i in inputs])  # all (H,W,C)
        if ntype == "iterable-array[bool]":
            return [SegmentationMapOnImage(attr_i, shape=image_i.shape)
                    for attr_i, image_i in zip(inputs, images)]
        return [SegmentationMapOnImage(
                    attr_i, shape=image_i.shape, nb_classes=1+np.max(attr_i))
                for attr_i, image_i in zip(inputs, images)]
    else:
        assert ntype == "iterable-SegmentationMapOnImage"
        return inputs  # len allowed to differ from len of images
<commit_msg>Fix expected shapes in normalize_segmentation_maps()<commit_after>from __future__ import print_function, division, absolute_import

import numpy as np

import imgaug.imgaug as ia
import imgaug.dtypes as iadt


def normalize_images(images):
    if images is None:
        return None
    elif ia.is_np_array(images):
        if images.ndim == 2:
            return images[np.newaxis, ..., np.newaxis]
        elif images.ndim == 3:
            return images[..., np.newaxis]
        else:
            return images
    elif ia.is_iterable(images):
        result = []
        for image in images:
            if image.ndim == 2:
                result.append(image[..., np.newaxis])
            else:
                assert image.ndim == 3
                result.append(image)
        return result
    raise ValueError(
        ("Expected argument 'images' to be any of the following: "
         + "None or array or iterable of array. Got type: %s.") % (
            type(images),)
    )


def normalize_heatmaps(inputs, images=None):
    # TODO get rid of this deferred import
    from imgaug.augmentables.heatmaps import HeatmapsOnImage

    ntype = estimate_heatmaps_norm_type(inputs)
    if ntype == "None":
        return None
    elif ntype == "array[float]":
        assert images is not None
        assert inputs.ndim == 4  # always (N,H,W,C)
        assert len(inputs) == len(images)
        return [HeatmapsOnImage(attr_i, shape=image_i.shape)
                for attr_i, image_i in zip(inputs, images)]
    elif ntype == "HeatmapsOnImage":
        return [inputs]
    elif ntype == "iterable[empty]":
        return None
    elif ntype == "iterable-array[float]":
        assert images is not None
        assert len(inputs) == len(images)
        assert all([attr_i.ndim == 3 for attr_i in inputs])  # all (H,W,C)
        return [HeatmapsOnImage(attr_i, shape=image_i.shape)
                for attr_i, image_i in zip(inputs, images)]
    else:
        assert ntype == "iterable-HeatmapsOnImage"
        return inputs  # len allowed to differ from len of images


def normalize_segmentation_maps(inputs, images=None):
    # TODO get rid of this deferred import
    from imgaug.augmentables.segmaps import SegmentationMapOnImage

    ntype = estimate_segmaps_norm_type(inputs)
    if ntype == "None":
        return None
    elif ntype in ["array[int]", "array[uint]", "array[bool]"]:
        assert images is not None
        assert inputs.ndim == 3  # always (N,H,W)
        assert len(inputs) == len(images)
        if ntype == "array[bool]":
            return [SegmentationMapOnImage(attr_i, shape=image_i.shape)
                    for attr_i, image_i in zip(inputs, images)]
        return [SegmentationMapOnImage(
                    attr_i, shape=image_i.shape, nb_classes=1+np.max(attr_i))
                for attr_i, image_i in zip(inputs, images)]
    elif ntype == "SegmentationMapOnImage":
        return [inputs]
    elif ntype == "iterable[empty]":
        return None
    elif ntype in ["iterable-array[int]",
                   "iterable-array[uint]",
                   "iterable-array[bool]"]:
        assert images is not None
        assert len(inputs) == len(images)
        assert all([attr_i.ndim == 2 for attr_i in inputs])  # all (H,W)
        if ntype == "iterable-array[bool]":
            return [SegmentationMapOnImage(attr_i, shape=image_i.shape)
                    for attr_i, image_i in zip(inputs, images)]
        return [SegmentationMapOnImage(
                    attr_i, shape=image_i.shape, nb_classes=1+np.max(attr_i))
                for attr_i, image_i in zip(inputs, images)]
    else:
        assert ntype == "iterable-SegmentationMapOnImage"
        return inputs  # len allowed to differ from len of images

{'fraction_non_alphanumeric': 0.10028003733831177, 'fraction_numerical': 0.0026670222696359514, 'mean_word_length': 2.290917068889864, 'pattern_counts': {'":': 20, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '231713', 'n_tokens_mistral': 2294, 'n_tokens_neox': 2212, 'n_words': 636}