Coverage for paintmaskanimator/auto_part_tracking.py: 91%
290 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 04:50 +0000
1"""確定した部位を保持し、線画の局所的な変形を使って再発見する。"""
2from collections import OrderedDict
3from dataclasses import dataclass
5import numpy as np
7from . import bucket_fill, line_flow, reference_check
10MEMORY_BLOCK = 128
11MAX_DONORS = 3
12MAX_FLOW_PAIRS = 64
13MAX_TRACK_STEPS = 12
14MAX_TILE_STEPS = 2
17@dataclass
18class _FlowPair:
19 forward: line_flow.Flow
20 backward: line_flow.Flow
21 x0: int
22 y0: int
24 def coordinates(self, ys, xs):
25 dy, dx = self.forward.at(ys - self.y0, xs - self.x0)
26 sy, sx = ys + dy, xs + dx
27 by, bx = self.backward.at(sy - self.y0, sx - self.x0)
28 return (np.rint(sy).astype(np.int64), np.rint(sx).astype(np.int64),
29 np.rint(sy + by).astype(np.int64), np.rint(sx + bx).astype(np.int64),
30 np.ones(ys.shape, dtype=bool))
33@dataclass
34class _FlowPath:
35 pairs: list
36 shape: tuple
38 def coordinates(self, ys, xs):
39 sy, sx = ys.astype(np.float64), xs.astype(np.float64)
40 valid = np.ones(ys.shape, dtype=bool)
41 height, width = self.shape
42 for pair in self.pairs:
43 dy, dx = pair.forward.at(sy - pair.y0, sx - pair.x0)
44 sy, sx = sy + dy, sx + dx
45 valid &= (sy >= 0) & (sy < height) & (sx >= 0) & (sx < width)
46 by, bx = sy.copy(), sx.copy()
47 for pair in reversed(self.pairs):
48 dy, dx = pair.backward.at(by - pair.y0, bx - pair.x0)
49 by, bx = by + dy, bx + dx
50 return (np.rint(sy).astype(np.int64), np.rint(sx).astype(np.int64),
51 np.rint(by).astype(np.int64), np.rint(bx).astype(np.int64), valid)
54def _in_mask(ys, xs, x0, y0, mask):
55 y, x = ys - y0, xs - x0
56 inside = (y >= 0) & (y < mask.shape[0]) & (x >= 0) & (x < mask.shape[1])
57 result = np.zeros(ys.shape, dtype=bool)
58 result[inside] = mask[y[inside], x[inside]]
59 return result
62def _near_black(cell, ys, xs):
63 height, width = cell.black.shape
64 result = np.zeros(ys.shape, dtype=bool)
65 for dy in (-1, 0, 1):
66 for dx in (-1, 0, 1):
67 result |= cell.black[np.clip(ys + dy, 0, height - 1), np.clip(xs + dx, 0, width - 1)]
68 return result
71class PartTracker:
72 """色を確定した場所の索引。途中の未塗りセルで古い確定部分を消さない。"""
74 def __init__(self, shape, geometry_for=None, alignment=None):
75 from .auto_alignment import AlignmentCache
76 self.shape = shape
77 self.geometry_for = geometry_for
78 self.alignment = alignment if alignment is not None else AlignmentCache(gpu="off")
79 self.coverage = {}
80 self.pairs = OrderedDict()
81 self.flow_solves = 0
82 self.matches = 0
83 self.long_matches = 0
85 def _coverage(self, bounds, mask):
86 x0, y0, x1, y1 = bounds
87 counts = {}
88 for ty in range(y0 // MEMORY_BLOCK, (y1 - 1) // MEMORY_BLOCK + 1):
89 for tx in range(x0 // MEMORY_BLOCK, (x1 - 1) // MEMORY_BLOCK + 1):
90 ax0, ay0 = max(x0, tx * MEMORY_BLOCK), max(y0, ty * MEMORY_BLOCK)
91 ax1, ay1 = min(x1, (tx + 1) * MEMORY_BLOCK), min(y1, (ty + 1) * MEMORY_BLOCK)
92 count = int(mask[ay0 - y0:ay1 - y0, ax0 - x0:ax1 - x0].sum())
93 if count >= 8:
94 counts[(ty, tx)] = count
95 return counts
97 def remember_fixed(self, index, guide):
98 known = (guide.values != 0) & ~guide.black & ~guide.primary & ~guide.open & ~guide.undecided
99 height, width = self.shape
100 self.coverage[index] = self._coverage((0, 0, width, height), known)
102 def remember(self, index, paint):
103 counts = {}
104 for _rgb, bounds, mask in paint.paints:
105 for key, count in self._coverage(bounds, mask).items():
106 counts[key] = min(MEMORY_BLOCK ** 2, counts.get(key, 0) + count)
107 # 更新するのはこのセルの確定部分だけ。他セルの最後の成功は保持する。
108 self.coverage[index] = counts
110 def donors(self, index, filled):
111 query = self._coverage(filled.bounds, filled.local)
112 total = sum(query.values())
113 if not total: 113 ↛ 114line 113 didn't jump to line 114 because the condition on line 113 was never true
114 return []
115 ranked = []
116 for other, known in self.coverage.items():
117 if other == index or abs(other - index) > MAX_TRACK_STEPS or not known:
118 continue
119 overlap = 0
120 for (ty, tx), count in query.items():
121 nearby = max((known.get((ty + dy, tx + dx), 0)
122 for dy in (-1, 0, 1) for dx in (-1, 0, 1)), default=0)
123 overlap += min(count, nearby)
124 if overlap >= max(8, total * 0.2): 124 ↛ 116line 124 didn't jump to line 116 because the condition on line 124 was always true
125 ranked.append((overlap / total, -abs(other - index), -other, other))
126 return [item[-1] for item in sorted(ranked, reverse=True)[:MAX_DONORS]]
128 def _pair(self, target, source):
129 from .auto_bucket import geometry_key
130 key = (geometry_key(target), geometry_key(source))
131 if key in self.pairs:
132 self.pairs.move_to_end(key)
133 return self.pairs[key]
134 reverse = key[1], key[0]
135 if reverse in self.pairs:
136 self.pairs.move_to_end(reverse)
137 found = self.pairs[reverse]
138 return _FlowPair(found.backward, found.forward, found.x0, found.y0)
139 if target.values.shape != self.shape or source.values.shape != self.shape: 139 ↛ 140line 139 didn't jump to line 140 because the condition on line 139 was never true
140 return None
141 first, second = bucket_fill.mask_bounds(target.boundary), bucket_fill.mask_bounds(source.boundary)
142 if first is None or second is None: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 return None
144 height, width = self.shape
145 x0, y0 = max(0, min(first[0], second[0]) - 32), max(0, min(first[1], second[1]) - 32)
146 x1, y1 = min(width, max(first[2], second[2]) + 32), min(height, max(first[3], second[3]) + 32)
147 left, right = target.boundary[y0:y1, x0:x1], source.boundary[y0:y1, x0:x1]
148 pair = _FlowPair(line_flow.estimate_flow(left, right), line_flow.estimate_flow(right, left),
149 x0, y0)
150 self.alignment.register_motion(target, source, pair.forward, x0, y0)
151 self.alignment.register_motion(source, target, pair.backward, x0, y0)
152 self.flow_solves += 1
153 self.pairs[key] = pair
154 if len(self.pairs) > MAX_FLOW_PAIRS: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true
155 self.pairs.popitem(last=False)
156 return pair
158 def _path(self, index, target, guide):
159 if self.geometry_for is None or guide.index is None:
160 return self._pair(target, guide.initial)
161 direction = 1 if guide.index > index else -1
162 pairs = []
163 for step in range(index, guide.index, direction):
164 pair = self._pair(self.geometry_for(step), self.geometry_for(step + direction))
165 if pair is None: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true
166 return None
167 pairs.append(pair)
168 return _FlowPath(pairs, self.shape)
170 def _color(self, index, target, filled, guide):
171 rgb = self._flow_color(index, target, filled, guide)
172 return rgb if rgb is not None else self._tiled_color(index, target, filled, guide)
174 def prime_motion(self, index, target, guides):
175 """通常の領域照合より先に、隣接セル間の共有変形場を用意する。"""
176 adjacent = [guide for guide in guides if guide.initial is not None
177 and guide.index is not None and abs(guide.index - index) == 1]
178 if adjacent:
179 self._path(index, target, adjacent[0])
181 def _flow_color(self, index, target, filled, guide):
182 from . import auto_paint as ap
183 source = guide.initial
184 if source is None:
185 return None
186 pair = self._path(index, target, guide)
187 if pair is None: 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true
188 return None
189 height, width = self.shape
190 x0, y0, x1, y1 = filled.bounds
191 ax0, ay0 = max(0, x0 - 1), max(0, y0 - 1)
192 ax1, ay1 = min(width, x1 + 1), min(height, y1 + 1)
193 local = np.zeros((ay1 - ay0, ax1 - ax0), dtype=bool)
194 local[y0 - ay0:y1 - ay0, x0 - ax0:x1 - ax0] = filled.local
195 outline = bucket_fill.dilate(local, 1) & ~local
196 outline &= target.boundary[ay0:ay1, ax0:ax1]
197 oy, ox = np.nonzero(outline)
198 if len(oy) < 8:
199 return None
200 oy, ox = oy + ay0, ox + ax0
201 sy, sx, by, bx, consistent = pair.coordinates(oy, ox)
202 # 長い無地の輪郭では接線方向の位置は決まらない。同じ部位の輪郭へ戻ればよい。
203 outline_band = bucket_fill.dilate(local, 1) & ~bucket_fill.erode(local, 2)
204 consistent &= _in_mask(by, bx, ax0, ay0, outline_band)
205 inside = (sy >= 0) & (sy < height) & (sx >= 0) & (sx < width)
206 sy, sx = np.clip(sy, 0, height - 1), np.clip(sx, 0, width - 1)
207 hits = inside & consistent & source.wide_boundary[sy, sx]
208 hits &= np.where(target.black[oy, ox], _near_black(source, sy, sx), source.wide_primary[sy, sx])
209 if hits.mean() < 0.5:
210 return None
211 quadrants = (oy >= (y0 + y1) / 2).astype(np.uint8) * 2 + (ox >= (x0 + x1) / 2)
212 for quadrant in range(4):
213 selected = quadrants == quadrant
214 if selected.sum() >= 4 and hits[selected].mean() < 0.35: 214 ↛ 215line 214 didn't jump to line 215 because the condition on line 214 was never true
215 return None # 一辺だけの一致を、その部位全体の対応と取り違えない。
216 ys, xs = np.nonzero(filled.local)
217 ys, xs = ys + y0, xs + x0
218 sy, sx, by, bx, consistent = pair.coordinates(ys, xs)
219 consistent &= _in_mask(by, bx, x0, y0, filled.local)
220 known = (sy >= 0) & (sy < height) & (sx >= 0) & (sx < width) & consistent
221 sy, sx = np.clip(sy, 0, height - 1), np.clip(sx, 0, width - 1)
222 known &= ~source.boundary[sy, sx] & ~guide.open[sy, sx] & ~guide.undecided[sy, sx]
223 confirmed = guide.confirmed
224 if confirmed is not None:
225 known &= confirmed[sy, sx]
226 known &= (guide.values[sy, sx] != 0) & ~guide.black[sy, sx] & ~guide.primary[sy, sx]
227 size = len(ys)
228 if known.sum() < size * 0.7:
229 return None
230 value, share = ap._majority(guide.values[sy[known], sx[known]])
231 if value is None or share < 0.95 or not ap._paintable(value):
232 return None
233 if size < 4096: 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 matched = known & (guide.values[sy, sx] == value)
235 labels = np.unique(guide.regions.labels[sy[matched], sx[matched]])
236 other = sum(int(guide.regions.sizes[label]) for label in labels)
237 if other > size * 4 or matched.sum() < other * 0.3:
238 return None
239 return reference_check._to_rgb(value)
241 def _tiled_color(self, index, target, filled, guide):
242 """滑らかな動きで合わない部位は、小区画を独立に追跡して色だけを集計する。"""
243 from . import auto_paint as ap
244 source = guide.initial
245 if source is None:
246 return None
247 if self.geometry_for is None or guide.index is None:
248 steps = [(target, source)]
249 else:
250 if abs(guide.index - index) > MAX_TILE_STEPS:
251 return None # 長い経路は共有した変形場で追い、同じFFT探索を区画ごとに繰り返さない。
252 direction = 1 if guide.index > index else -1
253 steps = [(self.geometry_for(step), self.geometry_for(step + direction))
254 for step in range(index, guide.index, direction)]
255 if any(left.values.shape != self.shape or right.values.shape != self.shape 255 ↛ 257line 255 didn't jump to line 257 because the condition on line 255 was never true
256 for left, right in steps):
257 return None
258 height, width = self.shape
259 x0, y0, x1, y1 = filled.bounds
260 outline = bucket_fill.dilate(filled.local, 1) & ~bucket_fill.erode(filled.local, 1)
261 samples, labels = [], set()
262 size = int(filled.local.sum())
263 tile_size = 64 if size < 262144 else 96
264 tiles = []
265 for ty in range(y0, y1, tile_size):
266 for tx in range(x0, x1, tile_size):
267 sub = filled.local[ty - y0:ty - y0 + tile_size, tx - x0:tx - x0 + tile_size]
268 if sub.sum() < 16: 268 ↛ 269line 268 didn't jump to line 269 because the condition on line 268 was never true
269 continue
270 oy, ox = np.nonzero(outline[ty - y0:ty - y0 + sub.shape[0],
271 tx - x0:tx - x0 + sub.shape[1]])
272 real = target.wide_boundary[oy + ty, ox + tx]
273 oy, ox = oy[real], ox[real]
274 if len(oy) < 4: 274 ↛ 275line 274 didn't jump to line 275 because the condition on line 274 was never true
275 continue
276 initial = (tx, ty, tx + sub.shape[1], ty + sub.shape[0])
277 tiles.append((tx, ty, sub, oy, ox, initial))
278 paths = self.alignment.track_many(steps, [tile[-1] for tile in tiles], self.shape)
279 for (tx, ty, sub, oy, ox, _initial), path in zip(tiles, paths):
280 if path is None:
281 continue
282 dy, dx, back_y, back_x = path
283 sy, sx = oy + ty + dy, ox + tx + dx
284 at = (sy >= 0) & (sy < height) & (sx >= 0) & (sx < width)
285 hit = np.zeros(at.shape, dtype=bool)
286 hit[at] = source.wide_boundary[sy[at], sx[at]]
287 syc, sxc = np.clip(sy, 0, height - 1), np.clip(sx, 0, width - 1)
288 hit &= np.where(_near_black(target, oy + ty, ox + tx),
289 _near_black(source, syc, sxc), source.wide_primary[syc, sxc])
290 if hit.mean() < 0.5:
291 continue
292 ys, xs = np.nonzero(sub)
293 ys, xs = ys + ty, xs + tx
294 sy, sx = ys + dy, xs + dx
295 known = (sy >= 0) & (sy < height) & (sx >= 0) & (sx < width)
296 sy, sx = np.clip(sy, 0, height - 1), np.clip(sx, 0, width - 1)
297 known &= _in_mask(sy + back_y, sx + back_x, x0, y0, filled.local)
298 known &= ~source.boundary[sy, sx] & ~guide.open[sy, sx] & ~guide.undecided[sy, sx]
299 known &= (guide.values[sy, sx] != 0) & ~guide.black[sy, sx] & ~guide.primary[sy, sx]
300 if guide.confirmed is not None:
301 known &= guide.confirmed[sy, sx]
302 if known.sum() < len(ys) * 0.5:
303 continue
304 values = guide.values[sy[known], sx[known]]
305 value, share = ap._majority(values)
306 if value is None or share < 0.9 or not ap._paintable(value):
307 continue
308 samples.append(values)
309 labels.update(np.unique(guide.regions.labels[sy[known], sx[known]]).tolist())
310 if not samples or sum(len(values) for values in samples) < size * 0.5:
311 return None
312 values = np.concatenate(samples)
313 value, share = ap._majority(values)
314 if value is None or share < 0.95 or not ap._paintable(value):
315 return None
316 if size < 4096 and sum(int(guide.regions.sizes[label]) for label in labels) > size * 4: 316 ↛ 317line 316 didn't jump to line 317 because the condition on line 316 was never true
317 return None
318 return reference_check._to_rgb(value)
320 def color(self, index, target, filled, guide_for):
321 if filled.target is not None or int(filled.local.sum()) < 16:
322 return None
323 # 外側背景を追跡対象にすると、全画像の空白を毎回処理することになる。
324 height, width = self.shape
325 if filled.bounds == (0, 0, width, height):
326 return None
327 answers = []
328 for other in self.donors(index, filled):
329 rgb = self._color(index, target, filled, guide_for(other))
330 if rgb is not None:
331 answers.append((other, rgb))
332 if not answers or any(rgb != answers[0][1] for _other, rgb in answers):
333 return None
334 self.matches += 1
335 self.long_matches += int(min(abs(other - index) for other, _rgb in answers) > 4)
336 return answers[0][1]