harmonization.recon

 1import numpy as np
 2
 3from itertools import product
 4from scipy.ndimage import zoom
 5
 6
 7def depimp_zoom(D, block_size, block_up, order=1):
 8
 9    block_size = np.array(block_size)
10    block_up = np.array(block_up)
11
12    if (len(block_up) == len(block_size)):
13        factor = block_size / block_up
14        zoomer = (*factor, 1)
15    else:
16        factor = block_size[:-1] / block_up
17        zoomer = (*factor, 1, 1)
18
19    # if we have a 4D block array and different last dimension, subsample it
20    # if (len(block_up) == len(block_size)) and (block_up[-1] != block_size[-1]):
21    size = tuple(block_up)
22
23    if (len(block_up) - 1) == len(block_size):
24        size = size + (block_size[-1],)
25
26    reshaped = D.reshape(size + (-1,))
27    D_depimpe = zoom(reshaped, zoomer, order=order).reshape(np.prod(block_size), -1)
28
29    return D_depimpe
30
31
32def reconstruct_from_blocks(patches, image_size, block_size, block_up, new_overlap, weights=None):
33
34    i_h, i_w, i_l = image_size[:3]
35
36    if len(patches.shape[1:]) == 3:
37        p_h, p_w, p_l = patches.shape[1:]
38    else:
39        p_h, p_w, p_l = patches.shape[1:-1]
40
41    img = np.zeros(image_size, dtype=np.float32)
42    img = np.pad(img, [(0, p_h), (0, p_w), (0, p_l), (0, 0)], 'constant', constant_values=(0, 0))
43    div = np.full(image_size, 1e-15, dtype=np.float32)
44    div = np.pad(div, [(0, p_h), (0, p_w), (0, p_l), (0, 0)], 'constant', constant_values=(1e-15, 1e-15))
45
46    # compute the dimensions of the patches array
47    n_h = i_h - p_h + 1
48    n_w = i_w - p_w + 1
49    n_l = i_l - p_l + 1
50
51    frac, nup = np.modf(np.divide(block_up, block_size[:-1]))
52    int_frac = np.ceil(frac).astype(np.int32)
53    nup = np.array(nup + int_frac, dtype=np.int32)
54
55    valid = (np.arange(0, block_up[0], nup[0]),
56             np.arange(0, block_up[1], nup[1]),
57             np.arange(0, block_up[2], nup[2]))
58
59    plus = [None, None, None]
60
61    for idx in range(3):
62        if len(valid[idx]) < block_size[idx]:
63            candidate = np.arange(block_up[idx])[(len(valid[idx]) - block_size[idx]):]
64
65            # if we already have filled the last index (like in 4/5), we have to find a valid location for it
66            while any(np.in1d(candidate, valid[idx])):
67                candidate -= 1
68
69            plus[idx] = candidate
70        else:
71            plus[idx] = ()
72
73    valid = (np.append(valid[0], plus[0]),
74             np.append(valid[1], plus[1]),
75             np.append(valid[2], plus[2]))
76
77    step = ([slice(i, i + p_h) for i in range(0, n_h + int_frac[0], new_overlap[0]) if i % block_up[0] in valid[0]],
78            [slice(j, j + p_w) for j in range(0, n_w + int_frac[1], new_overlap[1]) if j % block_up[1] in valid[1]],
79            [slice(k, k + p_l) for k in range(0, n_l + int_frac[2], new_overlap[2]) if k % block_up[2] in valid[2]])
80
81    ijk = list(product(*step))
82
83    for p, (i, j, k) in zip(patches, ijk):
84        img[i, j, k] += p
85        div[i, j, k] += 1
86
87    out = img[:-p_h, :-p_w, :-p_l] / div[:-p_h, :-p_w, :-p_l]
88    return out