<commit_before># Copyright 2021 <NAME>
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Union, List, Tuple
from abc import abstractmethod
import math
import warnings

import jax
import jax.numpy as jnp
from einops import rearrange

from .base import Transformation, BaseChain, InputType, same_type
from . import utils


class LazyCoordinates:
    _current_transform: jnp.ndarray = jnp.eye(3)
    _offsets: Union[jnp.ndarray, None] = None
    input_shape: Tuple[int, int]
    current_shape: Tuple[int, int]
    final_shape: Tuple[int, int]

    def __init__(self, shape: Tuple[int, int]):
        self.input_shape = shape
        self.current_shape = shape
        self.final_shape = shape

    def get_coordinate_grid(self) -> jnp.ndarray:
        H, W = self.final_shape
        coordinates = jnp.mgrid[0:H, 0:W] - jnp.array([H/2-0.5, W/2-0.5]).reshape(2, 1, 1)
        coordinates = utils.apply_perspective(coordinates, self._current_transform)

        if self._offsets is not None:
            coordinates = coordinates + self._offsets

        H, W = self.input_shape
        return coordinates + jnp.array([H/2-0.5, W/2-0.5]).reshape(2, 1, 1)

    def apply_to_points(self, points) -> jnp.ndarray:
        M_inv = jnp.linalg.inv(self._current_transform)

        H_in, W_in = self.input_shape
        H_out, W_out = self.final_shape
        c_x = jnp.array([H_in/2 - 0.5, W_in/2 - 0.5]).reshape(2, 1)
        c_y = jnp.array([H_out/2 - 0.5, W_out/2 - 0.5]).reshape(2, 1)
        points = points.T

        transformed_points = utils.apply_perspective(points - c_x, M_inv) + c_y
        if self._offsets is not None:
            # Need to do fix-point iteration
            points_iter = transformed_points
            offset_grid = rearrange(self._offsets, 'c h w -> h w c')
            for _ in range(2):
                # fix-point iteration
                offsets = utils.resample_image(offset_grid, points_iter, order=1).T
                points_iter = utils.apply_perspective(points - offsets - c_x, M_inv) + c_y
            transformed_points = points_iter

        return transformed_points.T

    def push_transform(self, M: jnp.ndarray):
        assert M.shape == (3, 3)
        self._current_transform = M @ self._current_transform
        self._dirty = True

    def apply_pixelwise_offsets(self, offsets):
        assert offsets.shape[1:] == self.final_shape
        if self._offsets == None:
            self._offsets = offsets
        else:
            self._offsets = self._offsets + offsets


class GeometricTransformation(Transformation):
    @abstractmethod
    def transform_coordinates(self, rng: jnp.ndarray, coordinates: LazyCoordinates, invert=False) -> LazyCoordinates:
        return coordinates

    def apply(self, rng: jnp.ndarray, inputs: jnp.ndarray, input_types: List[InputType]=None, invert=False) -> List[jnp.ndarray]:
        if input_types is None:
            input_types = self.input_types

        input_shape  = inputs[0].shape[:2]
        output_shape = self.output_shape(input_shape)
        if invert:
            if hasattr(self, 'shape_full'):
                output_shape = self.shape_full
            elif self.size_changing():
                raise ValueError("Can't invert a size-changing transformation without running it forward once.")
        else:
            self.shape_full = input_shape

        coordinates = LazyCoordinates(input_shape)
        coordinates.final_shape = output_shape

        if invert:
            coordinates.current_shape = output_shape

        self.transform_coordinates(rng, coordinates, invert)
        sampling_coords = coordinates.get_coordinate_grid()

        val = []
        for input, type in zip(inputs, input_types):
            current = None
            if same_type(type, InputType.IMAGE) or same_type(type, InputType.DENSE):
                # Linear Interpolation for Images
                current = utils.resample_image(input, sampling_coords, order=1)
            elif same_type(type, InputType.MASK):
                # Nearest Interpolation for Masks
                current = utils.resample_image(input, sampling_coords, order=0)
            elif same_type(type, InputType.KEYPOINTS):
                current = coordinates.apply_to_points(input)
            elif same_type(type, InputType.CONTOUR):
                current = coordinates.apply_to_points(input)
                current = jnp.where(jnp.linalg.det(coordinates._current_transform) < 0,
                    current[::-1],
                    current<commit_msg>Pad with 0 instead of edge vals<commit_after># Copyright 2021 <NAME>
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Union, List, Tuple
from abc import abstractmethod
import math
import warnings

import jax
import jax.numpy as jnp
from einops import rearrange

from .base import Transformation, BaseChain, InputType, same_type
from . import utils


class LazyCoordinates:
    _current_transform: jnp.ndarray = jnp.eye(3)
    _offsets: Union[jnp.ndarray, None] = None
    input_shape: Tuple[int, int]
    current_shape: Tuple[int, int]
    final_shape: Tuple[int, int]

    def __init__(self, shape: Tuple[int, int]):
        self.input_shape = shape
        self.current_shape = shape
        self.final_shape = shape

    def get_coordinate_grid(self) -> jnp.ndarray:
        H, W = self.final_shape
        coordinates = jnp.mgrid[0:H, 0:W] - jnp.array([H/2-0.5, W/2-0.5]).reshape(2, 1, 1)
        coordinates = utils.apply_perspective(coordinates, self._current_transform)

        if self._offsets is not None:
            coordinates = coordinates + self._offsets

        H, W = self.input_shape
        return coordinates + jnp.array([H/2-0.5, W/2-0.5]).reshape(2, 1, 1)

    def apply_to_points(self, points) -> jnp.ndarray:
        M_inv = jnp.linalg.inv(self._current_transform)

        H_in, W_in = self.input_shape
        H_out, W_out = self.final_shape
        c_x = jnp.array([H_in/2 - 0.5, W_in/2 - 0.5]).reshape(2, 1)
        c_y = jnp.array([H_out/2 - 0.5, W_out/2 - 0.5]).reshape(2, 1)
        points = points.T

        transformed_points = utils.apply_perspective(points - c_x, M_inv) + c_y
        if self._offsets is not None:
            # Need to do fix-point iteration
            points_iter = transformed_points
            offset_grid = rearrange(self._offsets, 'c h w -> h w c')
            for _ in range(2):
                # fix-point iteration
                offsets = utils.resample_image(offset_grid, points_iter, order=1).T
                points_iter = utils.apply_perspective(points - offsets - c_x, M_inv) + c_y
            transformed_points = points_iter

        return transformed_points.T

    def push_transform(self, M: jnp.ndarray):
        assert M.shape == (3, 3)
        self._current_transform = M @ self._current_transform
        self._dirty = True

    def apply_pixelwise_offsets(self, offsets):
        assert offsets.shape[1:] == self.final_shape
        if self._offsets == None:
            self._offsets = offsets
        else:
            self._offsets = self._offsets + offsets


class GeometricTransformation(Transformation):
    @abstractmethod
    def transform_coordinates(self, rng: jnp.ndarray, coordinates: LazyCoordinates, invert=False) -> LazyCoordinates:
        return coordinates

    def apply(self, rng: jnp.ndarray, inputs: jnp.ndarray, input_types: List[InputType]=None, invert=False) -> List[jnp.ndarray]:
        if input_types is None:
            input_types = self.input_types

        input_shape  = inputs[0].shape[:2]
        output_shape = self.output_shape(input_shape)
        if invert:
            if hasattr(self, 'shape_full'):
                output_shape = self.shape_full
            elif self.size_changing():
                raise ValueError("Can't invert a size-changing transformation without running it forward once.")
        else:
            self.shape_full = input_shape

        coordinates = LazyCoordinates(input_shape)
        coordinates.final_shape = output_shape

        if invert:
            coordinates.current_shape = output_shape

        self.transform_coordinates(rng, coordinates, invert)
        sampling_coords = coordinates.get_coordinate_grid()

        val = []
        for input, type in zip(inputs, input_types):
            current = None
            if same_type(type, InputType.IMAGE) or same_type(type, InputType.DENSE):
                # Linear Interpolation for Images
                current = utils.resample_image(input, sampling_coords, order=1, mode='constant')
            elif same_type(type, InputType.MASK):
                # Nearest Interpolation for Masks
                current = utils.resample_image(input, sampling_coords, order=0, mode='constant')
            elif same_type(type, InputType.KEYPOINTS):
                current = coordinates.apply_to_points(input)
            elif same_type(type, InputType.CONTOUR):
                current = coordinates.apply_to_points(input)
                current = jnp.where(jnp.linalg.det(coordinates._current_transform) < 0,
                    current[::-1],
                    current
{'fraction_non_alphanumeric': 0.0910524241229799, 'fraction_numerical': 0.011135199054000788, 'mean_word_length': 2.537469501568491, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '613970', 'n_tokens_mistral': 3099, 'n_tokens_neox': 2991, 'n_words': 795}