opt_flow_reg/stitcher.py at master · BayraktarLab/opt_flow_reg · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
from typing import List, Tuple, Union

import numpy as np

Image = np.ndarray


def get_slices(
    big_image: Image,
    hor_f: int,
    hor_t: int,
    ver_f: int,
    ver_t: int,
    padding: dict,
    overlap=0,
):
    y_axis = 0
    x_axis = 1
    # check if tile is over image boundary
    left_check = hor_f - padding["left"]
    top_check = ver_f - padding["top"]
    right_check = hor_t - big_image.shape[x_axis]
    bot_check = ver_t - big_image.shape[y_axis]

    left_pad_size = 0
    top_pad_size = 0
    right_pad_size = 0
    bot_pad_size = 0

    if left_check < 0:
        left_pad_size = abs(left_check)
        hor_f = 0
    if top_check < 0:
        top_pad_size = abs(top_check)
        ver_f = 0
    if right_check > 0:
        right_pad_size = right_check
        hor_t = big_image.shape[x_axis]
    if bot_check > 0:
        bot_pad_size = bot_check
        ver_t = big_image.shape[y_axis]

    big_image_slice = [slice(ver_f, ver_t), slice(hor_f, hor_t)]
    tile_shape = (ver_t - ver_f, hor_t - hor_f)
    tile_slice = [
        slice(top_pad_size + overlap, tile_shape[-2] + overlap),
        slice(left_pad_size + overlap, tile_shape[-1] + overlap),
    ]
    if len(big_image.shape) > 2:
        big_image_slice.append(slice(None))
        tile_slice.append(slice(None))
    return tuple(big_image_slice), tuple(tile_slice)


def stitch_image(img_list: List[Image], slicer_info: dict) -> Image:

    x_ntiles = slicer_info["ntiles"]["x"]
    y_ntiles = slicer_info["ntiles"]["y"]
    tile_shape = slicer_info["tile_shape"]
    overlap = slicer_info["overlap"]
    padding = slicer_info["padding"]

    x_axis = 1
    y_axis = 0
    ch_axis = 2

    tile_x_size = tile_shape[x_axis]
    tile_y_size = tile_shape[y_axis]

    big_image_x_size = (x_ntiles * tile_x_size) - padding["left"] - padding["right"]
    big_image_y_size = (y_ntiles * tile_y_size) - padding["top"] - padding["bottom"]

    if len(img_list[0].shape) == 2:
        big_image_shape = (big_image_y_size, big_image_x_size)
    elif len(img_list[0].shape) == 3:
        # must be a flow map
        big_image_shape = (big_image_y_size, big_image_x_size, 2)
    else:
        msg = f"Input image has unexpected dimensions - {str(img_list[0].shape)}"
        raise ValueError(msg)
    dtype = img_list[0].dtype
    big_image = np.zeros(big_image_shape, dtype=dtype)

    n = 0
    for i in range(0, y_ntiles):
        ver_f = i * tile_y_size
        ver_t = ver_f + tile_y_size

        for j in range(0, x_ntiles):
            hor_f = j * tile_x_size
            hor_t = hor_f + tile_x_size

            big_image_slice, tile_slice = get_slices(
                big_image, hor_f, hor_t, ver_f, ver_t, padding, overlap
            )
            tile = img_list[n]

            big_image[tuple(big_image_slice)] = tile[tuple(tile_slice)]
            n += 1

    return big_image