Spaces:
Runtime error
Runtime error
File size: 5,868 Bytes
c19ca42 |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Callable, List, Optional, Tuple
import numpy as np
from nodes.log import logger
from ...utils.utils import Padding, Region, Size, get_h_w_c
from ..image_utils import BorderType, create_border
def _pad_image(img: np.ndarray, min_size: Size):
h, w, _ = get_h_w_c(img)
min_w, min_h = min_size
x = max(0, min_w - w) / 2
y = max(0, min_h - h) / 2
padding = Padding(math.floor(y), math.floor(x), math.ceil(y), math.ceil(x))
return create_border(img, BorderType.REFLECT_MIRROR, padding), padding
@dataclass
class _Segment:
start: int
end: int
startPadding: int
endPadding: int
@property
def length(self) -> int:
return self.end - self.start
@property
def padded_length(self) -> int:
return self.end + self.endPadding - (self.start - self.startPadding)
def _exact_split_into_segments(length: int, exact: int, overlap: int) -> List[_Segment]:
"""
Splits the given length into segments of `exact` (padded) length.
Segments will overlap into each other with at least the given overlap.
"""
if length == exact:
# trivial
return [_Segment(0, exact, 0, 0)]
assert length > exact
assert exact > overlap * 2
result: List[_Segment] = []
def add(s: _Segment):
assert s.padded_length == exact
result.append(s)
# The current strategy is to go from left to right and to align segments
# such that we use the least overlap possible. The last segment will then
# be the smallest with potentially a lot of overlap.
# While this is easy to implement, it's actually not ideal. Ideally, we
# would want for the overlap to be distributed evenly between segments.
# However, this is complex to implement and the current method also works.
# we know that the first segment looks like this
add(_Segment(0, exact - overlap, 0, overlap))
while result[-1].end < length:
startPadding = overlap
start = result[-1].end
end = start + exact - overlap * 2
endPadding = overlap
if end + endPadding >= length:
# last segment
endPadding = 0
end = length
startPadding = exact - (end - start)
add(_Segment(start, end, startPadding, endPadding))
return result
def _exact_split_into_regions(
w: int,
h: int,
exact_w: int,
exact_h: int,
overlap: int,
) -> List[Tuple[Region, Padding]]:
"""
Returns a list of disjoint regions along with padding.
Each region plus its padding is guaranteed to have the given exact size.
The padding (if not zero) is guaranteed to be at least the given overlap value.
"""
# we can split x and y independently from each other and then combine the results
x_segments = _exact_split_into_segments(w, exact_w, overlap)
y_segments = _exact_split_into_segments(h, exact_h, overlap)
logger.info(f"chaiNNer: image is split into {len(x_segments)}x{len(y_segments)} tiles each exactly {exact_w}x{exact_h}px")
result: List[Tuple[Region, Padding]] = []
for y in y_segments:
for x in x_segments:
result.append(
(
Region(x.start, y.start, x.length, y.length),
Padding(y.startPadding, x.endPadding, y.endPadding, x.startPadding),
)
)
return result
def _exact_split_without_padding(
img: np.ndarray,
exact_size: Size,
upscale: Callable[[np.ndarray, Region], np.ndarray],
overlap: int,
) -> np.ndarray:
h, w, c = get_h_w_c(img)
exact_w, exact_h = exact_size
assert w >= exact_w and h >= exact_h
if (w, h) == exact_size:
return upscale(img, Region(0, 0, w, h))
# To allocate the result image, we need to know the upscale factor first,
# and we only get to know this factor after the first successful upscale.
result: Optional[np.ndarray] = None
scale: int = 0
regions = _exact_split_into_regions(w, h, exact_w, exact_h, overlap)
for tile, pad in regions:
padded_tile = tile.add_padding(pad)
upscale_result = upscale(padded_tile.read_from(img), padded_tile)
# figure out by how much the image was upscaled by
up_h, up_w, _ = get_h_w_c(upscale_result)
current_scale = up_h // padded_tile.height
assert current_scale > 0
assert padded_tile.height * current_scale == up_h
assert padded_tile.width * current_scale == up_w
if result is None:
# allocate the result image
scale = current_scale
result = np.zeros((h * scale, w * scale, c), dtype=np.float32)
assert current_scale == scale
# remove overlap padding
upscale_result = pad.scale(scale).remove_from(upscale_result)
# copy into result image
tile.scale(scale).write_into(result, upscale_result)
assert result is not None
# remove initially added padding
return result
def exact_split(
img: np.ndarray,
exact_size: Size,
upscale: Callable[[np.ndarray, Region], np.ndarray],
overlap: int = 16,
) -> np.ndarray:
"""
Splits the image into tiles with exactly the given tile size.
If the image is smaller than the given size, then it will be padded.
"""
# ensure that the image is at least as large as the given size
img, base_padding = _pad_image(img, exact_size)
h, w, _ = get_h_w_c(img)
result = _exact_split_without_padding(img, exact_size, upscale, overlap)
scale = get_h_w_c(result)[0] // h
if base_padding.empty:
return result
# remove initially added padding
return (
Region(0, 0, w, h).remove_padding(base_padding).scale(scale).read_from(result)
)
|