Coverage for paintmaskanimator/line_flow.py: 98%
154 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 04:50 +0000
1"""2枚の線画のあいだの動き（区画ごとのずれ）を、粗い解像度から細かい解像度へ求める。
3中割りのセル同士は、顔の向きや手の動きで場所ごとにずれ方が違い、大きく動く所もある。
4部分ごとに1つのずれを探すと、大きな部分では場所によって合わず、小さな部分では
5まわりの線が少なくて迷う。ここでは絵を ``BLOCK`` 画素の区画に分け、区画ごとのずれを
71. 縮小した線で広い範囲を探し（となりの区画の重なりも足して、ばらつきを抑える）、
82. 解像度を上げながら、そのずれのまわりだけを探し直す。
10線が無い区画は、となりの区画の答えに従う。重なりは「動かしたあとの手本の線
11（1画素太らせる）に、こちらの線がどれだけ重なったか」で数える。
12候補ごとの一致数は Rust の部品があればそちらで求め、無いときは numpy で求める。
13"""
14from dataclasses import dataclass
15from typing import Tuple
17import numpy as np
19from . import bucket_fill, native
21#: 区画の大きさ（画素）。
22BLOCK = 32
23#: いちばん粗い縮小率と、そこで探すずれの範囲（画素）。
24COARSE = 8
25REACH = 96
26#: 細かくするときの縮小率と、それぞれで探し直す範囲（その縮小率での画素）。
27REFINE_STEPS = (4, 2, 1)
28REFINE_REACH = 2
29#: となりの区画の重なりも足す重み。ずれがばらばらになるのを防ぐ。
30NEIGHBOUR_WEIGHT = 0.5
31#: 大きなずれほど少し割り引く（同じくらい重なるなら、動きの小さいほうを選ぶ）。
32MOTION_PENALTY = 0.002
33#: 区画の線がこれより少なければ、その区画は「手がかりなし」。
34MIN_LINE_PIXELS = 6
37@dataclass
38class Flow:
39 """区画ごとのずれ ``(dy, dx)``（画素）と、その確かさ（0〜1、重なった線の割合）。"""
41 dy: np.ndarray
42 dx: np.ndarray
43 score: np.ndarray
44 #: 区画にある、こちらの線の画素数（手がかりの量）。
45 lines: np.ndarray
46 shape: Tuple[int, int]
48 def at(self, ys, xs):
49 """画素 ``(ys, xs)`` のずれ。区画の中心のあいだを直線でつないで滑らかにする。"""
50 ys, xs = np.broadcast_arrays(ys, xs)
51 dy = np.empty(ys.shape, dtype=np.float64)
52 dx = np.empty(xs.shape, dtype=np.float64)
53 # 大きな背景領域でも一時配列をキャッシュに収め、縦横の補間位置を共有する。
54 flat_y, flat_x = ys.ravel(), xs.ravel()
55 out_y, out_x = dy.ravel(), dx.ravel()
56 for start in range(0, ys.size, 65536):
57 end = start + 65536
58 weights = _bilinear_weights(self.dy.shape, flat_y[start:end], flat_x[start:end])
59 out_y[start:end] = _interpolate(self.dy, weights)
60 out_x[start:end] = _interpolate(self.dx, weights)
61 return dy, dx
63 def score_at(self, ys, xs):
64 by = np.clip(ys // BLOCK, 0, self.score.shape[0] - 1)
65 bx = np.clip(xs // BLOCK, 0, self.score.shape[1] - 1)
66 return self.score[by, bx]
69def _bilinear(grid, ys, xs):
70 """区画の中心に置いた値 ``grid`` を、画素 ``(ys, xs)`` で直線補間する。"""
71 return _interpolate(grid, _bilinear_weights(grid.shape, ys, xs))
74def _bilinear_weights(shape, ys, xs):
75 rows, cols = shape
76 fy = np.clip((ys + 0.5) / BLOCK - 0.5, 0, rows - 1)
77 fx = np.clip((xs + 0.5) / BLOCK - 0.5, 0, cols - 1)
78 y0 = np.floor(fy).astype(np.int64)
79 x0 = np.floor(fx).astype(np.int64)
80 y1 = np.minimum(y0 + 1, rows - 1)
81 x1 = np.minimum(x0 + 1, cols - 1)
82 wy = fy - y0
83 wx = fx - x0
84 return y0, x0, y1, x1, wy, wx
87def _interpolate(grid, weights):
88 y0, x0, y1, x1, wy, wx = weights
89 top = grid[y0, x0] * (1 - wx) + grid[y0, x1] * wx
90 bottom = grid[y1, x0] * (1 - wx) + grid[y1, x1] * wx
91 return top * (1 - wy) + bottom * wy
94def _pool(mask, scale, rows, cols):
95 """``scale`` 画素四方に1画素でも線があれば線とする縮小。区画の数に合わせて余白を足す。"""
96 height, width = rows * BLOCK // scale, cols * BLOCK // scale
97 padded = np.zeros((rows * BLOCK, cols * BLOCK), dtype=bool)
98 padded[:mask.shape[0], :mask.shape[1]] = mask
99 if scale == 1:
100 return padded
101 return padded.reshape(height, scale, width, scale).any(axis=(1, 3))
104def _block_sum(values, rows, cols):
105 size = values.shape[0] // rows
106 return values.reshape(rows, size, cols, size).sum(axis=(1, 3))
109def _with_neighbours(scores):
110 """区画ごとの重なりに、上下左右ななめの区画の重なりを ``NEIGHBOUR_WEIGHT`` 倍して足す。"""
111 padded = np.pad(scores, [(0, 0)] * (scores.ndim - 2) + [(1, 1), (1, 1)])
112 total = np.zeros_like(scores, dtype=np.float64)
113 rows, cols = scores.shape[-2:]
114 for dy in (0, 1, 2):
115 for dx in (0, 1, 2):
116 weight = 1.0 if dy == 1 and dx == 1 else NEIGHBOUR_WEIGHT
117 total += weight * padded[..., dy:dy + rows, dx:dx + cols]
118 return total
121def _shifted(mask, dy, dx):
122 """``out[y, x] = mask[y + dy, x + dx]``（はみ出しは偽）。"""
123 out = np.zeros_like(mask)
124 height, width = mask.shape
125 ys0, ys1 = max(0, -dy), min(height, height - dy)
126 xs0, xs1 = max(0, -dx), min(width, width - dx)
127 if ys0 < ys1 and xs0 < xs1:
128 out[ys0:ys1, xs0:xs1] = mask[ys0 + dy:ys1 + dy, xs0 + dx:xs1 + dx]
129 return out
132def _scores_numpy(target, guide, rows, cols, reach, base_y=None, base_x=None):
133 """候補ごとの区画の一致画素数（Rust と同じ定義の参照実装）。"""
134 shifts = [(dy, dx) for dy in range(-reach, reach + 1)
135 for dx in range(-reach, reach + 1)]
136 scores = np.empty((len(shifts), rows, cols))
137 if base_y is None or base_x is None:
138 for index, (dy, dx) in enumerate(shifts):
139 scores[index] = _block_sum(target & _shifted(guide, dy, dx), rows, cols)
140 return scores
141 size = target.shape[0] // rows
142 height, width = target.shape
143 ys, xs = np.mgrid[0:height, 0:width]
144 base_y = np.repeat(np.repeat(base_y, size, 0), size, 1)
145 base_x = np.repeat(np.repeat(base_x, size, 0), size, 1)
146 for index, (dy, dx) in enumerate(shifts):
147 gy, gx = ys + base_y + dy, xs + base_x + dx
148 inside = (gy >= 0) & (gy < height) & (gx >= 0) & (gx < width)
149 hit = np.zeros_like(target)
150 hit[inside] = guide[gy[inside], gx[inside]]
151 scores[index] = _block_sum(target & hit, rows, cols)
152 return scores
155def _scores(target, guide, rows, cols, reach, base_y=None, base_x=None):
156 zero = np.zeros((rows, cols), dtype=np.int64) if base_y is None else None
157 return native.line_flow_scores(
158 target, guide, target.shape[0] // rows,
159 zero if base_y is None else base_y, zero if base_x is None else base_x, reach,
160 )
163def estimate_flow(target_lines, guide_lines) -> Flow:
164 """``target_lines`` の各区画が、``guide_lines`` のどこにあるか（真偽マスク同士）。"""
165 height, width = target_lines.shape
166 rows, cols = -(-height // BLOCK), -(-width // BLOCK)
168 # 1. 粗い解像度で広く探す。
169 scale = COARSE
170 target = _pool(target_lines, scale, rows, cols)
171 guide = bucket_fill.dilate(_pool(guide_lines, scale, rows, cols), 1)
172 counts = _block_sum(target, rows, cols).astype(np.float64)
173 reach = REACH // scale
174 shifts = [(dy, dx) for dy in range(-reach, reach + 1) for dx in range(-reach, reach + 1)]
175 scores = _scores(target, guide, rows, cols, reach)
176 penalty = np.array([MOTION_PENALTY * scale * np.hypot(dy, dx) for dy, dx in shifts])
177 ratio = _with_neighbours(scores) / np.maximum(_with_neighbours(counts[None]), 1.0)
178 best = np.argmax(ratio - penalty[:, None, None], axis=0)
179 flow_y = np.array([shifts[i][0] for i in best.ravel()]).reshape(rows, cols) * scale
180 flow_x = np.array([shifts[i][1] for i in best.ravel()]).reshape(rows, cols) * scale
182 # 2. 解像度を上げて、今のずれのまわりだけを探し直す。
183 score = np.zeros((rows, cols))
184 for scale in REFINE_STEPS:
185 target = _pool(target_lines, scale, rows, cols)
186 guide = bucket_fill.dilate(_pool(guide_lines, scale, rows, cols), 1)
187 counts = _block_sum(target, rows, cols).astype(np.float64)
188 base_y = np.round(flow_y / scale).astype(np.int64)
189 base_x = np.round(flow_x / scale).astype(np.int64)
190 deltas = [(dy, dx) for dy in range(-REFINE_REACH, REFINE_REACH + 1)
191 for dx in range(-REFINE_REACH, REFINE_REACH + 1)]
192 scores = _scores(target, guide, rows, cols, REFINE_REACH, base_y, base_x)
193 own = scores / np.maximum(counts, 1.0)
194 ratio = _with_neighbours(scores) / np.maximum(_with_neighbours(counts[None]), 1.0)
195 penalty = np.array([MOTION_PENALTY * scale * np.hypot(dy, dx) for dy, dx in deltas])
196 best = np.argmax(ratio - penalty[:, None, None], axis=0)
197 flow_y = flow_y + np.array([deltas[i][0] for i in best.ravel()]).reshape(rows, cols) * scale
198 flow_x = flow_x + np.array([deltas[i][1] for i in best.ravel()]).reshape(rows, cols) * scale
199 score = np.take_along_axis(own, best[None], axis=0)[0]
200 lines = counts
202 # 線の無い区画は、まわりの区画のずれ（線の多さで重み付け）に従う。
203 lines = _block_sum(_pool(target_lines, 1, rows, cols), rows, cols)
204 weak = lines < MIN_LINE_PIXELS
205 if weak.any() and (~weak).any():
206 weight = np.where(weak, 0.0, lines.astype(np.float64))
207 for _ in range(max(rows, cols)): 207 ↛ 220line 207 didn't jump to line 220 because the loop on line 207 didn't complete
208 if not weak.any():
209 break
210 spread = _with_neighbours(weight[None])[0]
211 fill_y = _with_neighbours((weight * flow_y)[None])[0]
212 fill_x = _with_neighbours((weight * flow_x)[None])[0]
213 reached = weak & (spread > 0)
214 if not reached.any(): 214 ↛ 215line 214 didn't jump to line 215 because the condition on line 214 was never true
215 break
216 flow_y = np.where(reached, fill_y / np.maximum(spread, 1e-9), flow_y)
217 flow_x = np.where(reached, fill_x / np.maximum(spread, 1e-9), flow_x)
218 weight = np.where(reached, 1.0, weight)
219 weak = weak & ~reached
220 score = np.where(lines < MIN_LINE_PIXELS, 0.0, score)
221 return Flow(flow_y.astype(np.float64), flow_x.astype(np.float64), score, lines,
222 (height, width))
225__all__ = ["BLOCK", "Flow", "estimate_flow"]