<commit_before>"""
take in a numpy tensor, and reshape last 2 dimensions as images to fit the desired shape
"""
import multiprocessing as mp
import numpy as np
import skimage.io
import skimage.transform
from custom_warnings import deprecated


tform_identity = skimage.transform.AffineTransform()
NO_AUGMENT_PARAMS = {
    "zoom_x": 1.0,
    "zoom_y": 1.0,
    "rotate": 0.0,
    "shear": 0.0,
    "skew_x": 0.0,
    "skew_y": 0.0,
    "translate_x": 0.0,
    "flip_vert": 0.0,
    "roll_time": 0.0,
    "flip_time": 0.0,
}

def resize_to_make_it_fit(images, output_shape=(50, 50)):
    max_time = max(images[i].shape[0] for i in xrange(len(images)))
    final_shape = (len(images),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    volume_change = []
    #result.reshape((final_shape[0],-1) + output_shape)
    for i, mri_slice in enumerate(images):
        mri_slice = mri_slice.reshape((-1,)+mri_slice.shape[-2:])
        scaling = max(mri_slice[0].shape[-2]/output_shape[-2], mri_slice[0].shape[-1]/output_shape[-1])
        tform = build_rescale_transform(scaling, mri_slice[0].shape[-2:], target_shape=output_shape)

        for j, frame in enumerate(mri_slice):
            # TODO: can't this be done better?
            result[i,j] = fast_warp(frame, tform, output_shape=output_shape)

        A = tform.params[:2, :2]
        volume_change.append(np.linalg.norm(A[:,0]) * np.linalg.norm(A[:,1]))
        assert tform.params[2,2] == 1, (tform.params[2,2],)

    #result.reshape(final_shape)
    return result, volume_change


def normscale_resize_and_augment(slices, output_shape=(50, 50), augment=None,
                                 pixel_spacing=(1,1), shift_center=(.4, .5),
                                 normalised_patch_size=(200,200)):
    """Normalizes the scale, augments, and crops the image.
    """
    if not pixel_spacing[0] == pixel_spacing[1]:
        raise NotImplementedError("Only supports square pixels")

    # No augmentation:
    if augment is None:
        augment = NO_AUGMENT_PARAMS

    current_shape = slices[0].shape[-2:]
    normalised_shape = tuple(int(float(d)*ps) for d,ps in zip(current_shape, pixel_spacing))

    max_time = max(slices[i].shape[0] for i in xrange(len(slices)))
    final_shape = (len(slices),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    for i, mri_slice in enumerate(slices):
        # For each slice, build a transformation that extracts the right patch,
        # and augments the data.
        # First, we scale the images such that they all have the same scale
        norm_rescaling = 1./pixel_spacing[0]
        tform_normscale = build_rescale_transform(
            norm_rescaling, mri_slice[0].shape[-2:], target_shape=normalised_shape)
        # Next, we shift the center of the image to the left (assumes upside_up normalisation)
        tform_shift_center, tform_shift_uncenter = (
            build_shift_center_transform(
                normalised_shape, shift_center, normalised_patch_size))

        # zooming is OK
        augment_tform = build_augmentation_transform(**augment)

        patch_scale = max(
            normalised_patch_size[0]/output_shape[0],
            normalised_patch_size[1]/output_shape[1])
        tform_patch_scale = build_rescale_transform(
            patch_scale, normalised_patch_size, target_shape=output_shape)

        # x and y axis transform
        total_tform = tform_patch_scale + tform_shift_uncenter + augment_tform + tform_shift_center + tform_normscale

        # Time axis transform
        t_map = range(mri_slice.shape[0])
        if "roll_time" in augment:
            t_map = np.roll(t_map, int(np.floor(augment["roll_time"])))
        if "flip_time" in augment and augment["flip_time"] > 0.5:
            t_map = t_map[::-1]

        for j, frame in enumerate(mri_slice):
            j_shifted = t_map[j]
            result[i,j_shifted] = fast_warp(frame, total_tform, output_shape=output_shape)

    return result


NRMSC_DEFAULT_SHIFT_CENTER = (.4, .5)
def normscale_resize_and_augment_2(slices, output_shape=(50, 50), augment=None,
                                   pixel_spacing=(1,1), shift_center=(None, None),
                                   normalised_patch_size=(200,200)):
    """Normalizes the scale, augments, and crops the image.
    Fixed a bug !
    """
    if not pixel_spacing[0] == pixel_spacing[1]:
        raise NotImplementedError("Only supports square pixels")

    if shift_center == (None, None):
        shift_center = NRMSC_DEFAULT_SHIFT_CENTER

    # No augmentation:
    if augment is None:
        augment = NO_AUGMENT_PARAMS

    current_shape = slices[0].shape[-2:]
    normalised_shape = tuple(int(float(d)*ps) for d,ps in zip(current_shape, pixel_spacing))

    max_time = max(slices[i].shape[0] for i in xrange(len(slices)))
    final_shape = (len(slices),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    for i, mri_slice in enumerate(slices):
        # For each slice, build a transformation that extracts the right patch,
        # and augments the data.
        # First, we scale the images such that they all have the same scale
        norm_rescaling = 1./pixel_spacing[0]
        tform_normscale = build_rescale_transform(
            norm_rescaling, mri_slice[0].shape[-2:], target_shape=normalised_shape)
        # Next, we shift the center of the image to the left (assumes upside_up normalisation)
        tform_shift_center, tform_shift_uncenter = (
            build_shift_center_transform(
                normalised_shape, shift_center, normalised_patch_size))

        # No zooming to preserve scale!
        augment["zoom_x"] = 1.0
        augment["zoom_y"] = 1.0<commit_msg>FUCK GIT<commit_after>"""
take in a numpy tensor, and reshape last 2 dimensions as images to fit the desired shape
"""
import multiprocessing as mp
import numpy as np
import skimage.io
import skimage.transform
from custom_warnings import deprecated


tform_identity = skimage.transform.AffineTransform()
NO_AUGMENT_PARAMS = {
    "zoom_x": 1.0,
    "zoom_y": 1.0,
    "rotate": 0.0,
    "shear": 0.0,
    "skew_x": 0.0,
    "skew_y": 0.0,
    "translate_x": 0.0,
    "flip_vert": 0.0,
    "roll_time": 0.0,
    "flip_time": 0.0,
}

def resize_to_make_it_fit(images, output_shape=(50, 50)):
    max_time = max(images[i].shape[0] for i in xrange(len(images)))
    final_shape = (len(images),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    volume_change = []
    #result.reshape((final_shape[0],-1) + output_shape)
    for i, mri_slice in enumerate(images):
        mri_slice = mri_slice.reshape((-1,)+mri_slice.shape[-2:])
        scaling = max(mri_slice[0].shape[-2]/output_shape[-2], mri_slice[0].shape[-1]/output_shape[-1])
        tform = build_rescale_transform(scaling, mri_slice[0].shape[-2:], target_shape=output_shape)

        for j, frame in enumerate(mri_slice):
            # TODO: can't this be done better?
            result[i,j] = fast_warp(frame, tform, output_shape=output_shape)

        A = tform.params[:2, :2]
        volume_change.append(np.linalg.norm(A[:,0]) * np.linalg.norm(A[:,1]))
        assert tform.params[2,2] == 1, (tform.params[2,2],)

    #result.reshape(final_shape)
    return result, volume_change


def normscale_resize_and_augment(slices, output_shape=(50, 50), augment=None,
                                 pixel_spacing=(1,1), shift_center=(.4, .5),
                                 normalised_patch_size=(200,200)):
    """Normalizes the scale, augments, and crops the image.
    """
    if not pixel_spacing[0] == pixel_spacing[1]:
        raise NotImplementedError("Only supports square pixels")

    # No augmentation:
    if augment is None:
        augment = NO_AUGMENT_PARAMS

    current_shape = slices[0].shape[-2:]
    normalised_shape = tuple(int(float(d)*ps) for d,ps in zip(current_shape, pixel_spacing))

    max_time = max(slices[i].shape[0] for i in xrange(len(slices)))
    final_shape = (len(slices),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    for i, mri_slice in enumerate(slices):
        # For each slice, build a transformation that extracts the right patch,
        # and augments the data.
        # First, we scale the images such that they all have the same scale
        norm_rescaling = 1./pixel_spacing[0]
        tform_normscale = build_rescale_transform(
            norm_rescaling, mri_slice[0].shape[-2:], target_shape=normalised_shape)
        # Next, we shift the center of the image to the left (assumes upside_up normalisation)
        tform_shift_center, tform_shift_uncenter = (
            build_shift_center_transform(
                normalised_shape, shift_center, normalised_patch_size))

        # zooming is OK
        augment_tform = build_augmentation_transform(**augment)

        patch_scale = max(
            normalised_patch_size[0]/output_shape[0],
            normalised_patch_size[1]/output_shape[1])
        tform_patch_scale = build_rescale_transform(
            patch_scale, normalised_patch_size, target_shape=output_shape)

        # x and y axis transform
        total_tform = tform_patch_scale + tform_shift_uncenter + augment_tform + tform_shift_center + tform_normscale

        # Time axis transform
        t_map = range(mri_slice.shape[0])
        if "roll_time" in augment:
            t_map = np.roll(t_map, int(np.floor(augment["roll_time"])))
        if "flip_time" in augment and augment["flip_time"] > 0.5:
            t_map = t_map[::-1]

        for j, frame in enumerate(mri_slice):
            j_shifted = t_map[j]
            result[i,j_shifted] = fast_warp(frame, total_tform, output_shape=output_shape)

    return result


NRMSC_DEFAULT_SHIFT_CENTER = (.4, .5)
def normscale_resize_and_augment_2(slices, output_shape=(50, 50), augment=None,
                                   pixel_spacing=(1,1), shift_center=(None, None),
                                   normalised_patch_size=(200,200)):
    """Normalizes the scale, augments, and crops the image.
    Fixed a bug !
    """
    if not pixel_spacing[0] == pixel_spacing[1]:
        raise NotImplementedError("Only supports square pixels")

    if shift_center == (None, None):
        shift_center = NRMSC_DEFAULT_SHIFT_CENTER

    # No augmentation:
    if augment is None:
        augment = NO_AUGMENT_PARAMS

    current_shape = slices[0].shape[-2:]
    normalised_shape = tuple(int(float(d)*ps) for d,ps in zip(current_shape, pixel_spacing))

    max_time = max(slices[i].shape[0] for i in xrange(len(slices)))
    final_shape = (len(slices),max_time) + output_shape
    result = np.zeros(final_shape, dtype="float32")

    for i, mri_slice in enumerate(slices):
        # For each slice, build a transformation that extracts the right patch,
        # and augments the data.
        # First, we scale the images such that they all have the same scale
        norm_rescaling = 1./pixel_spacing[0]
        tform_normscale = build_rescale_transform(
            norm_rescaling, mri_slice[0].shape[-2:], target_shape=normalised_shape)
        # Next, we shift the center of the image to the left (assumes upside_up normalisation)
        tform_shift_center, tform_shift_uncenter = (
            build_shift_center_transform(
                normalised_shape, shift_center, normalised_patch_size))

{'fraction_non_alphanumeric': 0.10884234510559986, 'fraction_numerical': 0.019279642450267285, 'mean_word_length': 3.058321479374111, 'pattern_counts': {'":': 20, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6842624', 'n_tokens_mistral': 3891, 'n_tokens_neox': 3696, 'n_words': 847}