import importlib
import math
import os
from pathlib import Path
from typing import Optional, Tuple

import cv2
import numpy as np
import requests
import torch

import kornia as K


def read_img_from_url(url: str, resize_to: Optional[Tuple[int, int]] = None) -> torch.Tensor:
    # perform request
    response = requests.get(url).content
    # convert to array of ints
    nparr = np.frombuffer(response, np.uint8)
    # convert to image array and resize
    img: np.ndarray = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)[..., :3]
    # convert the image to a tensor
    img_t: torch.Tensor = K.utils.image_to_tensor(img, keepdim=False)  # 1xCxHXW
    img_t = img_t.float() / 255.0
    if resize_to is None:
        img_t = K.geometry.resize(img_t, 184)
    else:
        img_t = K.geometry.resize(img_t, resize_to)
    return img_t


def main():
    # load the images
    BASE_IMAGE_URL1: str = "https://raw.githubusercontent.com/kornia/data/main/panda.jpg"  # augmentation
    BASE_IMAGE_URL2: str = "https://raw.githubusercontent.com/kornia/data/main/simba.png"  # color
    BASE_IMAGE_URL3: str = "https://raw.githubusercontent.com/kornia/data/main/girona.png"  # enhance
    BASE_IMAGE_URL4: str = "https://raw.githubusercontent.com/kornia/data/main/baby_giraffe.png"  # morphology
    BASE_IMAGE_URL5: str = "https://raw.githubusercontent.com/kornia/data/main/persistencia_memoria.jpg"  # filters
    BASE_IMAGE_URL6: str = "https://raw.githubusercontent.com/kornia/data/main/delorean.png"  # geometry
    OUTPUT_PATH = Path(__file__).absolute().parent / "source/_static/img"

    os.makedirs(OUTPUT_PATH, exist_ok=True)
    print(f"Pointing images to path {OUTPUT_PATH}.")
    img1 = read_img_from_url(BASE_IMAGE_URL1)
    img2 = read_img_from_url(BASE_IMAGE_URL2, img1.shape[-2:])
    img3 = read_img_from_url(BASE_IMAGE_URL3, img1.shape[-2:])
    img4 = read_img_from_url(BASE_IMAGE_URL4)
    img5 = read_img_from_url(BASE_IMAGE_URL5, (234, 320))
    img6 = read_img_from_url(BASE_IMAGE_URL6)

    # TODO: make this more generic for modules out of kornia.augmentation
    # Dictionary containing the transforms to generate the sample images:
    # Key: Name of the transform class.
    # Value: (parameters, num_samples, seed)
    mod = importlib.import_module("kornia.augmentation")
    augmentations_list: dict = {
        "CenterCrop": ((184, 184), 1, 2018),
        "ColorJitter": ((0.3, 0.3, 0.3, 0.3), 2, 2018),
        "RandomAffine": (((-15.0, 20.0), (0.1, 0.1), (0.7, 1.3), 20), 2, 2019),
        "RandomBoxBlur": (((7, 7),), 1, 2020),
        "RandomCrop": ((img1.shape[-2:], (50, 50)), 2, 2020),
        "RandomChannelShuffle": ((), 1, 2020),
        "RandomElasticTransform": (((63, 63), (32, 32), (2.0, 2.0)), 2, 2018),
        "RandomEqualize": ((), 1, 2020),
        "RandomErasing": (((0.2, 0.4), (0.3, 1 / 0.3)), 2, 2017),
        "RandomFisheye": ((torch.tensor([-0.3, 0.3]), torch.tensor([-0.3, 0.3]), torch.tensor([0.9, 1.0])), 2, 2020),
        "RandomGaussianBlur": (((3, 3), (0.1, 2.0)), 1, 2020),
        "RandomGaussianNoise": ((0.0, 0.05), 1, 2020),
        "RandomGrayscale": ((), 1, 2020),
        "RandomHorizontalFlip": ((), 1, 2020),
        "RandomInvert": ((), 1, 2020),
        "RandomMotionBlur": ((7, 35.0, 0.5), 2, 2020),
        "RandomPerspective": ((0.2,), 2, 2020),
        "RandomPlanckianJitter": ((), 2, 2022),
        "RandomPosterize": (((1, 4),), 2, 2016),
        "RandomResizedCrop": ((img1.shape[-2:], (1.0, 2.0), (1.0, 2.0)), 2, 2020),
        "RandomRotation": ((45.0,), 2, 2019),
        "RandomSharpness": ((16.0,), 1, 2019),
        "RandomSolarize": ((0.2, 0.2), 2, 2019),
        "RandomVerticalFlip": ((), 1, 2020),
        "RandomThinPlateSpline": ((), 1, 2020),
    }

    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, num_samples, seed) in augmentations_list.items():
        img_in = img1.repeat(num_samples, 1, 1, 1)
        # dynamically create the class instance
        cls = getattr(mod, aug_name)
        aug = cls(*args, p=1.0)
        # set seed
        torch.manual_seed(seed)
        # apply the augmentaiton to the image and concat
        out = aug(img_in)

        if aug_name == "CenterCrop":
            h, w = img1.shape[-2:]
            h_new, w_new = out.shape[-2:]
            h_dif, w_dif = int(h - h_new), int(w - w_new)
            out = torch.nn.functional.pad(out, (w_dif // 2, w_dif // 2, 0, h_dif))

        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
        sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
        print(f"Generated image example for {aug_name}. {sig}")

    mod = importlib.import_module("kornia.augmentation")
    mix_augmentations_list: dict = {
        "RandomMixUp": (((0.3, 0.4),), 2, 20),
        "RandomCutMix": ((img1.shape[-2], img1.shape[-1]), 2, 2019),
    }
    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, num_samples, seed) in mix_augmentations_list.items():
        img_in = torch.cat([img1, img2])
        # dynamically create the class instance
        cls = getattr(mod, aug_name)
        aug = cls(*args, p=1.0)
        # set seed
        torch.manual_seed(seed)
        # apply the augmentaiton to the image and concat
        out, _ = aug(img_in, torch.tensor([0, 1]))
        out = torch.cat([img_in[0], img_in[1], *(out[i] for i in range(out.size(0)))], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
        sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
        print(f"Generated image example for {aug_name}. {sig}")

    mod = importlib.import_module("kornia.color")
    color_transforms_list: dict = {
        "grayscale_to_rgb": ((), 3),
        "rgb_to_bgr": ((), 1),
        "rgb_to_grayscale": ((), 1),
        "rgb_to_hsv": ((), 1),
        "rgb_to_hls": ((), 1),
        "rgb_to_luv": ((), 1),
        "rgb_to_lab": ((), 1),
        # "rgb_to_rgba": ((1.,), 1),
        "rgb_to_xyz": ((), 1),
        "rgb_to_ycbcr": ((), 1),
        "rgb_to_yuv": ((), 1),
        "rgb_to_linear_rgb": ((), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in color_transforms_list.items():
        # import function and apply
        fn = getattr(mod, fn_name)
        if fn_name == "grayscale_to_rgb":
            out = fn(K.color.rgb_to_grayscale(img2), *args)
        else:
            out = fn(img2, *args)
        # perform normalization to visualize
        if fn_name == "rgb_to_lab":
            out = out[:, :1] / 100.0
        elif fn_name == "rgb_to_hsv":
            out[:, :1] = out[:, :1] / 2 * math.pi
        elif fn_name == "rgb_to_luv":
            out = out[:, :1] / 116.0
        # repeat channels for grayscale
        if out.shape[1] != 3:
            out = out.repeat(1, 3, 1, 1)
        # save the output image
        if fn_name == "grayscale_to_rgb":
            out = torch.cat(
                [K.color.rgb_to_grayscale(img2[0]).repeat(3, 1, 1), *(out[i] for i in range(out.size(0)))], dim=-1
            )
        else:
            out = torch.cat([img2[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.enhance module
    mod = importlib.import_module("kornia.enhance")
    transforms: dict = {
        "adjust_brightness": ((torch.tensor([0.25, 0.5]),), 2),
        "adjust_contrast": ((torch.tensor([0.65, 0.5]),), 2),
        "adjust_gamma": ((torch.tensor([0.85, 0.75]), 2.0), 2),
        "adjust_hue": ((torch.tensor([-math.pi / 4, math.pi / 4]),), 2),
        "adjust_saturation": ((torch.tensor([1.0, 2.0]),), 2),
        "solarize": ((torch.tensor([0.8, 0.5]), torch.tensor([-0.25, 0.25])), 2),
        "posterize": ((torch.tensor([4, 2]),), 2),
        "sharpness": ((torch.tensor([1.0, 2.5]),), 2),
        "equalize": ((), 1),
        "invert": ((), 1),
        "equalize_clahe": ((), 1),
        "add_weighted": ((0.75, 0.25, 2.0), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img3.repeat(num_samples, 1, 1, 1)
        if fn_name == "add_weighted":
            args_in = (img_in, args[0], img2, args[1], args[2])
        else:
            args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.morphology module
    mod = importlib.import_module("kornia.morphology")
    kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
    transforms: dict = {
        "dilation": ((kernel,), 1),
        "erosion": ((kernel,), 1),
        "opening": ((kernel,), 1),
        "closing": ((kernel,), 1),
        "gradient": ((kernel,), 1),
        "top_hat": ((kernel,), 1),
        "bottom_hat": ((kernel,), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img4.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        # import pdb;pdb.set_trace()
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.filters module
    mod = importlib.import_module("kornia.filters")
    kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
    transforms: dict = {
        "box_blur": (((5, 5),), 1),
        "median_blur": (((5, 5),), 1),
        "gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
        "motion_blur": ((5, 90.0, 1.0), 1),
        "max_blur_pool2d": ((5,), 1),
        "blur_pool2d": ((5,), 1),
        "unsharp_mask": (((5, 5), (1.5, 1.5)), 1),
        "laplacian": ((5,), 1),
        "sobel": ((), 1),
        "spatial_gradient": ((), 1),
        "canny": ((), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img5.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        if fn_name in ("max_blur_pool2d", "blur_pool2d"):
            out = K.geometry.resize(out, img_in.shape[-2:])
        if fn_name == "canny":
            out = out[1].repeat(1, 3, 1, 1)
        if isinstance(out, torch.Tensor):
            out = out.clamp(min=0.0, max=1.0)
        if fn_name in ("laplacian", "sobel", "spatial_gradient", "canny"):
            out = K.enhance.normalize_min_max(out)
        if fn_name == "spatial_gradient":
            out = out.permute(2, 1, 0, 3, 4).squeeze()
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.geometry.transform module
    mod = importlib.import_module("kornia.geometry.transform")
    h, w = img6.shape[-2:]

    def _get_tps_args():
        src = torch.tensor([[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0], [0.0, 0.0]]]).repeat(2, 1, 1)  # Bx5x2
        dst = src + torch.distributions.Uniform(-0.2, 0.2).rsample((2, 5, 2))
        kernel, affine = K.geometry.transform.get_tps_transform(dst, src)
        return src, kernel, affine

    transforms: dict = {
        "warp_affine": (
            (
                K.geometry.transform.get_affine_matrix2d(
                    translations=torch.zeros(2, 2),
                    center=(torch.tensor([w, h]) / 2).repeat(2, 1),
                    scale=torch.distributions.Uniform(0.5, 1.5).rsample((2, 2)),
                    angle=torch.tensor([-25.0, 25.0]),
                )[:, :2, :3],
                (h, w),
            ),
            2,
        ),
        "remap": (
            (
                *(K.utils.create_meshgrid(h, w, normalized_coordinates=True) - 0.25).unbind(-1),
                'bilinear',
                'zeros',
                True,
                True,
            ),
            1,
        ),
        "warp_image_tps": ((_get_tps_args()), 2),
        "rotate": ((torch.tensor([-15.0, 25.0]),), 2),
        "translate": ((torch.tensor([[10.0, -15], [50.0, -25.0]]),), 2),
        "scale": ((torch.tensor([[0.5, 1.25], [1.0, 1.5]]),), 2),
        "shear": ((torch.tensor([[0.1, -0.2], [-0.2, 0.1]]),), 2),
        "rot180": ((), 1),
        "hflip": ((), 1),
        "vflip": ((), 1),
        "resize": (((120, 220),), 1),
        "rescale": ((0.5,), 1),
        "elastic_transform2d": ((torch.rand(1, 2, h, w) * 2 - 1, (63, 63), (32, 32), (4.0, 4.0)), 1),
        "pyrdown": ((), 1),
        "pyrup": ((), 1),
        "build_pyramid": ((3,), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img6.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        if fn_name in ("resize", "rescale", "pyrdown", "pyrup"):
            h_new, w_new = out.shape[-2:]
            out = torch.nn.functional.pad(out, (0, (w - w_new), 0, (h - h_new)))
        if fn_name == "build_pyramid":
            _out = []
            for pyr in out[1:]:
                h_new, w_new = pyr.shape[-2:]
                out_tmp = torch.nn.functional.pad(pyr, (0, (w - w_new), 0, (h - h_new)))
                _out.append(out_tmp)
            out = torch.cat(_out)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")


if __name__ == "__main__":
    main()

[((19, 12, 19, 45), 'numpy.frombuffer', 'np.frombuffer', ({(19, 26, 19, 34): 'response', (19, 36, 19, 44): 'np.uint8'}, {}), '(response, np.uint8)', True, 'import numpy as np\n'), ((23, 26, 23, 69), 'kornia.utils.image_to_tensor', 'K.utils.image_to_tensor', (), '', True, 'import kornia as K\n'), ((42, 4, 42, 43), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((55, 10, 55, 56), 'importlib.import_module', 'importlib.import_module', ({(55, 34, 55, 55): '"""kornia.augmentation"""'}, {}), "('kornia.augmentation')", False, 'import importlib\n'), ((108, 10, 108, 56), 'importlib.import_module', 'importlib.import_module', ({(108, 34, 108, 55): '"""kornia.augmentation"""'}, {}), "('kornia.augmentation')", False, 'import importlib\n'), ((130, 10, 130, 49), 'importlib.import_module', 'importlib.import_module', ({(130, 34, 130, 48): '"""kornia.color"""'}, {}), "('kornia.color')", False, 'import importlib\n'), ((176, 10, 176, 51), 'importlib.import_module', 'importlib.import_module', ({(176, 34, 176, 50): '"""kornia.enhance"""'}, {}), "('kornia.enhance')", False, 'import importlib\n'), ((209, 10, 209, 54), 'importlib.import_module', 'importlib.import_module', ({(209, 34, 209, 53): '"""kornia.morphology"""'}, {}), "('kornia.morphology')", False, 'import importlib\n'), ((210, 13, 210, 60), 'torch.tensor', 'torch.tensor', ({(210, 26, 210, 59): '[[0, 1, 0], [1, 1, 1], [0, 1, 0]]'}, {}), '([[0, 1, 0], [1, 1, 1], [0, 1, 0]])', False, 'import torch\n'), ((236, 10, 236, 51), 'importlib.import_module', 'importlib.import_module', ({(236, 34, 236, 50): '"""kornia.filters"""'}, {}), "('kornia.filters')", False, 'import importlib\n'), ((237, 13, 237, 60), 'torch.tensor', 'torch.tensor', ({(237, 26, 237, 59): '[[0, 1, 0], [1, 1, 1], [0, 1, 0]]'}, {}), '([[0, 1, 0], [1, 1, 1], [0, 1, 0]])', False, 'import torch\n'), ((276, 10, 276, 62), 'importlib.import_module', 'importlib.import_module', ({(276, 34, 276, 61): '"""kornia.geometry.transform"""'}, {}), "('kornia.geometry.transform')", False, 'import importlib\n'), ((17, 15, 17, 32), 'requests.get', 'requests.get', ({(17, 28, 17, 31): 'url'}, {}), '(url)', False, 'import requests\n'), ((21, 22, 21, 63), 'cv2.imdecode', 'cv2.imdecode', ({(21, 35, 21, 40): 'nparr', (21, 42, 21, 62): 'cv2.IMREAD_UNCHANGED'}, {}), '(nparr, cv2.IMREAD_UNCHANGED)', False, 'import cv2\n'), ((26, 16, 26, 45), 'kornia.geometry.resize', 'K.geometry.resize', ({(26, 34, 26, 39): 'img_t', (26, 41, 26, 44): '184'}, {}), '(img_t, 184)', True, 'import kornia as K\n'), ((28, 16, 28, 51), 'kornia.geometry.resize', 'K.geometry.resize', ({(28, 34, 28, 39): 'img_t', (28, 41, 28, 50): 'resize_to'}, {}), '(img_t, resize_to)', True, 'import kornia as K\n'), ((91, 8, 91, 31), 'torch.manual_seed', 'torch.manual_seed', ({(91, 26, 91, 30): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((115, 17, 115, 40), 'torch.cat', 'torch.cat', ({(115, 27, 115, 39): '[img1, img2]'}, {}), '([img1, img2])', False, 'import torch\n'), ((120, 8, 120, 31), 'torch.manual_seed', 'torch.manual_seed', ({(120, 26, 120, 30): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((282, 25, 282, 73), 'kornia.geometry.transform.get_tps_transform', 'K.geometry.transform.get_tps_transform', ({(282, 64, 282, 67): 'dst', (282, 69, 282, 72): 'src'}, {}), '(dst, src)', True, 'import kornia as K\n'), ((99, 18, 99, 82), 'torch.nn.functional.pad', 'torch.nn.functional.pad', ({(99, 42, 99, 45): 'out', (99, 47, 99, 81): '(w_dif // 2, w_dif // 2, 0, h_dif)'}, {}), '(out, (w_dif // 2, w_dif // 2, 0, h_dif))', False, 'import torch\n'), ((122, 29, 122, 49), 'torch.tensor', 'torch.tensor', ({(122, 42, 122, 48): '[0, 1]'}, {}), '([0, 1])', False, 'import torch\n'), ((259, 18, 259, 59), 'kornia.geometry.resize', 'K.geometry.resize', ({(259, 36, 259, 39): 'out', (259, 41, 259, 58): 'img_in.shape[-2:]'}, {}), '(out, img_in.shape[-2:])', True, 'import kornia as K\n'), ((265, 18, 265, 50), 'kornia.enhance.normalize_min_max', 'K.enhance.normalize_min_max', ({(265, 46, 265, 49): 'out'}, {}), '(out)', True, 'import kornia as K\n'), ((332, 18, 332, 80), 'torch.nn.functional.pad', 'torch.nn.functional.pad', ({(332, 42, 332, 45): 'out', (332, 47, 332, 79): '(0, w - w_new, 0, h - h_new)'}, {}), '(out, (0, w - w_new, 0, h - h_new))', False, 'import torch\n'), ((339, 18, 339, 33), 'torch.cat', 'torch.cat', ({(339, 28, 339, 32): '_out'}, {}), '(_out)', False, 'import torch\n'), ((66, 27, 66, 52), 'torch.tensor', 'torch.tensor', ({(66, 40, 66, 51): '[-0.3, 0.3]'}, {}), '([-0.3, 0.3])', False, 'import torch\n'), ((66, 54, 66, 79), 'torch.tensor', 'torch.tensor', ({(66, 67, 66, 78): '[-0.3, 0.3]'}, {}), '([-0.3, 0.3])', False, 'import torch\n'), ((66, 81, 66, 105), 'torch.tensor', 'torch.tensor', ({(66, 94, 66, 104): '[0.9, 1.0]'}, {}), '([0.9, 1.0])', False, 'import torch\n'), ((150, 21, 150, 51), 'kornia.color.rgb_to_grayscale', 'K.color.rgb_to_grayscale', ({(150, 46, 150, 50): 'img2'}, {}), '(img2)', True, 'import kornia as K\n'), ((178, 31, 178, 56), 'torch.tensor', 'torch.tensor', ({(178, 44, 178, 55): '[0.25, 0.5]'}, {}), '([0.25, 0.5])', False, 'import torch\n'), ((179, 29, 179, 54), 'torch.tensor', 'torch.tensor', ({(179, 42, 179, 53): '[0.65, 0.5]'}, {}), '([0.65, 0.5])', False, 'import torch\n'), ((180, 26, 180, 52), 'torch.tensor', 'torch.tensor', ({(180, 39, 180, 51): '[0.85, 0.75]'}, {}), '([0.85, 0.75])', False, 'import torch\n'), ((181, 24, 181, 65), 'torch.tensor', 'torch.tensor', ({(181, 37, 181, 64): '[-math.pi / 4, math.pi / 4]'}, {}), '([-math.pi / 4, math.pi / 4])', False, 'import torch\n'), ((182, 31, 182, 55), 'torch.tensor', 'torch.tensor', ({(182, 44, 182, 54): '[1.0, 2.0]'}, {}), '([1.0, 2.0])', False, 'import torch\n'), ((183, 22, 183, 46), 'torch.tensor', 'torch.tensor', ({(183, 35, 183, 45): '[0.8, 0.5]'}, {}), '([0.8, 0.5])', False, 'import torch\n'), ((183, 48, 183, 75), 'torch.tensor', 'torch.tensor', ({(183, 61, 183, 74): '[-0.25, 0.25]'}, {}), '([-0.25, 0.25])', False, 'import torch\n'), ((184, 23, 184, 43), 'torch.tensor', 'torch.tensor', ({(184, 36, 184, 42): '[4, 2]'}, {}), '([4, 2])', False, 'import torch\n'), ((185, 23, 185, 47), 'torch.tensor', 'torch.tensor', ({(185, 36, 185, 46): '[1.0, 2.5]'}, {}), '([1.0, 2.5])', False, 'import torch\n'), ((280, 14, 280, 95), 'torch.tensor', 'torch.tensor', ({(280, 27, 280, 94): '[[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0], [0.0, 0.0]]]'}, {}), '([[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0], [0.0, \n    0.0]]])', False, 'import torch\n'), ((309, 20, 309, 47), 'torch.tensor', 'torch.tensor', ({(309, 33, 309, 46): '[-15.0, 25.0]'}, {}), '([-15.0, 25.0])', False, 'import torch\n'), ((310, 23, 310, 65), 'torch.tensor', 'torch.tensor', ({(310, 36, 310, 64): '[[10.0, -15], [50.0, -25.0]]'}, {}), '([[10.0, -15], [50.0, -25.0]])', False, 'import torch\n'), ((311, 19, 311, 58), 'torch.tensor', 'torch.tensor', ({(311, 32, 311, 57): '[[0.5, 1.25], [1.0, 1.5]]'}, {}), '([[0.5, 1.25], [1.0, 1.5]])', False, 'import torch\n'), ((312, 19, 312, 59), 'torch.tensor', 'torch.tensor', ({(312, 32, 312, 58): '[[0.1, -0.2], [-0.2, 0.1]]'}, {}), '([[0.1, -0.2], [-0.2, 0.1]])', False, 'import torch\n'), ((337, 26, 337, 88), 'torch.nn.functional.pad', 'torch.nn.functional.pad', ({(337, 50, 337, 53): 'pyr', (337, 55, 337, 87): '(0, w - w_new, 0, h - h_new)'}, {}), '(pyr, (0, w - w_new, 0, h - h_new))', False, 'import torch\n'), ((40, 18, 40, 32), 'pathlib.Path', 'Path', ({(40, 23, 40, 31): '__file__'}, {}), '(__file__)', False, 'from pathlib import Path\n'), ((281, 20, 281, 58), 'torch.distributions.Uniform', 'torch.distributions.Uniform', ({(281, 48, 281, 52): '(-0.2)', (281, 54, 281, 57): '(0.2)'}, {}), '(-0.2, 0.2)', False, 'import torch\n'), ((318, 33, 318, 55), 'torch.rand', 'torch.rand', ({(318, 44, 318, 45): '(1)', (318, 47, 318, 48): '(2)', (318, 50, 318, 51): 'h', (318, 53, 318, 54): 'w'}, {}), '(1, 2, h, w)', False, 'import torch\n'), ((166, 17, 166, 50), 'kornia.color.rgb_to_grayscale', 'K.color.rgb_to_grayscale', ({(166, 42, 166, 49): 'img2[0]'}, {}), '(img2[0])', True, 'import kornia as K\n'), ((289, 33, 289, 50), 'torch.zeros', 'torch.zeros', ({(289, 45, 289, 46): '(2)', (289, 48, 289, 49): '(2)'}, {}), '(2, 2)', False, 'import torch\n'), ((292, 26, 292, 53), 'torch.tensor', 'torch.tensor', ({(292, 39, 292, 52): '[-25.0, 25.0]'}, {}), '([-25.0, 25.0])', False, 'import torch\n'), ((300, 18, 300, 76), 'kornia.utils.create_meshgrid', 'K.utils.create_meshgrid', (), '', True, 'import kornia as K\n'), ((291, 26, 291, 63), 'torch.distributions.Uniform', 'torch.distributions.Uniform', ({(291, 54, 291, 57): '(0.5)', (291, 59, 291, 62): '(1.5)'}, {}), '(0.5, 1.5)', False, 'import torch\n'), ((290, 28, 290, 48), 'torch.tensor', 'torch.tensor', ({(290, 41, 290, 47): '[w, h]'}, {}), '([w, h])', False, 'import torch\n')]