<commit_before>""" Contains Batch classes for images """

import os   # pylint: disable=unused-import
import traceback

try:
    import blosc   # pylint: disable=unused-import
except ImportError:
    pass
import numpy as np
try:
    import PIL.Image
except ImportError:
    pass
try:
    import scipy.ndimage
except ImportError:
    pass
try:
    from numba import njit
except ImportError:
    pass

from ..batch import Batch
from ..decorators import action, inbatch_parallel, any_action_failed



CROP_CENTER = -1
CROP_00 = -2


@njit(nogil=True)
def random_crop_numba(images, shape):
    """ Fill-in new_images with random crops from images """
    new_images = np.zeros((images.shape[0],) + shape, dtype=images.dtype)
    if images.shape[2] - shape[0] > 0:
        origin_x = np.random.randint(0, images.shape[2] - shape[0], size=images.shape[0])
    else:
        origin_x = np.zeros(images.shape[0], dtype=np.array(images.shape).dtype)
    if images.shape[1] - shape[1] > 0:
        origin_y = np.random.randint(0, images.shape[1] - shape[1], size=images.shape[0])
    else:
        origin_y = np.zeros(images.shape[0], dtype=np.array(images.shape).dtype)
    for i in range(images.shape[0]):
        x = slice(origin_x[i], origin_x[i] + shape[0])
        y = slice(origin_y[i], origin_y[i] + shape[1])
        new_images[i, :, :] = images[i, y, x]
    return new_images

@njit(nogil=True)
def calc_origin(image_coord, shape_coord, crop):
    """ Return origin for preserve_shape """
    if crop == CROP_00:
        origin = 0
    elif crop == CROP_CENTER:
        origin = np.abs(shape_coord - image_coord) // 2
    return origin

@njit(nogil=True)
def calc_coords(image_coord, shape_coord, crop):
    """ Return coords for preserve_shape """
    if image_coord < shape_coord:
        new_image_origin = calc_origin(image_coord, shape_coord, crop)
        image_origin = 0
        image_len = image_coord
    else:
        new_image_origin = 0
        image_origin = calc_origin(image_coord, shape_coord, crop)
        image_len = shape_coord
    return image_origin, new_image_origin, image_len

@njit(nogil=True)
def preserve_shape_numba(image_shape, shape, crop=CROP_CENTER):
    """ Change the image shape by cropping and adding empty pixels to fit the given shape """
    x, new_x, len_x = calc_coords(image_shape[0], shape[0], crop)
    y, new_y, len_y = calc_coords(image_shape[1], shape[1], crop)
    new_x = new_x, new_x + len_x
    x = x, x + len_x
    new_y = new_y, new_y + len_y
    y = y, y + len_y

    return new_x, new_y, x, y



class BasicImagesBatch(Batch):
    """ Batch class for 2D images """
    @property
    def components(self):
        return "images", "labels"

    def assemble(self, all_res, *args, **kwargs):
        """ Assemble the batch after a parallel action """
        _ = all_res, args, kwargs
        raise NotImplementedError("Use ImagesBatch")

    def _assemble_load(self, all_res, *args, **kwargs):
        """ Build the batch data after loading data from files """
        _ = all_res, args, kwargs
        return self

    @action
    def load(self, src, fmt=None, components=None, *args, **kwargs):
        """ Load data """
        return super().load(src, fmt, components, *args, **kwargs)

    @action
    def dump(self, dst, fmt=None, components=None, *args, **kwargs):
        """ Save data to a file or a memory object """
        return super().dump(dst, fmt, components=None, *args, **kwargs)

    @action
    def noop(self):
        """ Do nothing """
        return self<commit_msg>Changed components definition<commit_after>""" Contains Batch classes for images """

import os   # pylint: disable=unused-import
import traceback

try:
    import blosc   # pylint: disable=unused-import
except ImportError:
    pass
import numpy as np
try:
    import PIL.Image
except ImportError:
    pass
try:
    import scipy.ndimage
except ImportError:
    pass
try:
    from numba import njit
except ImportError:
    pass

from ..batch import Batch
from ..decorators import action, inbatch_parallel, any_action_failed



CROP_CENTER = -1
CROP_00 = -2


@njit(nogil=True)
def random_crop_numba(images, shape):
    """ Fill-in new_images with random crops from images """
    new_images = np.zeros((images.shape[0],) + shape, dtype=images.dtype)
    if images.shape[2] - shape[0] > 0:
        origin_x = np.random.randint(0, images.shape[2] - shape[0], size=images.shape[0])
    else:
        origin_x = np.zeros(images.shape[0], dtype=np.array(images.shape).dtype)
    if images.shape[1] - shape[1] > 0:
        origin_y = np.random.randint(0, images.shape[1] - shape[1], size=images.shape[0])
    else:
        origin_y = np.zeros(images.shape[0], dtype=np.array(images.shape).dtype)
    for i in range(images.shape[0]):
        x = slice(origin_x[i], origin_x[i] + shape[0])
        y = slice(origin_y[i], origin_y[i] + shape[1])
        new_images[i, :, :] = images[i, y, x]
    return new_images

@njit(nogil=True)
def calc_origin(image_coord, shape_coord, crop):
    """ Return origin for preserve_shape """
    if crop == CROP_00:
        origin = 0
    elif crop == CROP_CENTER:
        origin = np.abs(shape_coord - image_coord) // 2
    return origin

@njit(nogil=True)
def calc_coords(image_coord, shape_coord, crop):
    """ Return coords for preserve_shape """
    if image_coord < shape_coord:
        new_image_origin = calc_origin(image_coord, shape_coord, crop)
        image_origin = 0
        image_len = image_coord
    else:
        new_image_origin = 0
        image_origin = calc_origin(image_coord, shape_coord, crop)
        image_len = shape_coord
    return image_origin, new_image_origin, image_len

@njit(nogil=True)
def preserve_shape_numba(image_shape, shape, crop=CROP_CENTER):
    """ Change the image shape by cropping and adding empty pixels to fit the given shape """
    x, new_x, len_x = calc_coords(image_shape[0], shape[0], crop)
    y, new_y, len_y = calc_coords(image_shape[1], shape[1], crop)
    new_x = new_x, new_x + len_x
    x = x, x + len_x
    new_y = new_y, new_y + len_y
    y = y, y + len_y

    return new_x, new_y, x, y



class BasicImagesBatch(Batch):
    """ Batch class for 2D images """
    components = "images", "labels"

    def assemble(self, all_res, *args, **kwargs):
        """ Assemble the batch after a parallel action """
        _ = all_res, args, kwargs
        raise NotImplementedError("Use ImagesBatch")

    def _assemble_load(self, all_res, *args, **kwargs):
        """ Build the batch data after loading data from files """
        _ = all_res, args, kwargs
        return self

    @action
    def load(self, src, fmt=None, components=None, *args, **kwargs):
        """ Load data """
        return super().load(src, fmt, components, *args, **kwargs)

    @action
    def dump(self, dst, fmt=None, components=None, *args, **kwargs):
        """ Save data to a file or a memory object """
        return super().dump(dst, fmt, components=None, *args, **kwargs)

    @action
    def noop(self):
        """ Do nothing """
        return self
{'fraction_non_alphanumeric': 0.12234794908062235, 'fraction_numerical': 0.009900990099009901, 'mean_word_length': 3.1520845566647093, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '2239781', 'n_tokens_mistral': 2405, 'n_tokens_neox': 2376, 'n_words': 642}