test_superresolution.py 2.09 KB
Newer Older
Stephan Seitz's avatar
Stephan Seitz committed
1
2
3
4
5
6
7
8
9
# -*- coding: utf-8 -*-
#
# Copyright © 2019 Stephan Seitz <stephan.seitz@fau.de>
#
# Distributed under terms of the GPLv3 license.

"""

"""
10
11
from os.path import dirname, join

Stephan Seitz's avatar
Stephan Seitz committed
12
import numpy as np
13
import skimage.io
Stephan Seitz's avatar
Stephan Seitz committed
14
import sympy
Stephan Seitz's avatar
Stephan Seitz committed
15
16

import pystencils
Stephan Seitz's avatar
Stephan Seitz committed
17
from pystencils_reco.resampling import downsample, resample, scale_transform, translate
Stephan Seitz's avatar
Stephan Seitz committed
18
19
20
21
22
23
24
25
26
27

try:
    import pyconrad.autoinit
except Exception:
    import unittest.mock
    pyconrad = unittest.mock.MagicMock()


def test_superresolution():

28
    x, y = np.random.rand(20, 10), np.zeros((20, 10))
Stephan Seitz's avatar
Stephan Seitz committed
29
30
31
32
33
34
35
36
37
38
39

    kernel = scale_transform(x, y, 0.5).compile()
    print(pystencils.show_code(kernel))
    kernel()

    pyconrad.show_everything()


def test_downsample():
    shape = (20, 10)

40
    x, y = np.random.rand(*shape), np.zeros(tuple(s // 2 for s in shape))
Stephan Seitz's avatar
Stephan Seitz committed
41
42
43
44
45
46

    kernel = downsample(x, y, 2).compile()
    print(pystencils.show_code(kernel))
    kernel()

    pyconrad.show_everything()
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65


def test_warp():
    import torch
    NUM_LENNAS = 5
    perturbation = 0.1

    lenna_file = join(dirname(__file__), "test_data", "lenna.png")
    lenna = skimage.io.imread(lenna_file, as_gray=True).astype(np.float32)

    warp_vectors = list(perturbation * torch.randn(lenna.shape + (2,)) for _ in range(NUM_LENNAS))

    warped = [torch.zeros(lenna.shape) for _ in range(NUM_LENNAS)]

    warp_kernel = translate(lenna, warped[0], pystencils.autodiff.ArrayWrapper(
        warp_vectors[0], index_dimensions=1), interpolation_mode='linear').compile()

    for i in range(len(warped)):
        warp_kernel(lenna[i], warped[i], warp_vectors[i])
Stephan Seitz's avatar
Stephan Seitz committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82


def test_polar_transform():
    x, y = pystencils.fields('x, y:  float32[2d]')

    x.set_coordinate_origin_to_field_center()
    y.set_coordinate_origin_to_field_center()
    y.coordinate_transform = lambda x: sympy.Matrix((x.norm(), sympy.atan2(*x) / (2 * sympy.pi) * y.shape[1]))

    lenna_file = join(dirname(__file__), "test_data", "lenna.png")
    lenna = skimage.io.imread(lenna_file, as_gray=True).astype(np.float32)

    transformed = np.zeros((500, 500), np.float32)

    resample(x, y).compile()(x=lenna, y=transformed)

    pyconrad.imshow(transformed)