Coverage for paintmaskanimator/auto_alignment.py: 74%
424 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 04:50 +0000
1"""位置合わせを依存段階ごとにまとめ、結果を再利用する。"""
2from collections import OrderedDict, defaultdict
3from concurrent.futures import ThreadPoolExecutor
4from contextlib import nullcontext
5from dataclasses import dataclass
6from hashlib import blake2b
7import logging
8import os
9from time import perf_counter
10from typing import Any, Callable, TypeVar
12import numpy as np
14_Answer = TypeVar("_Answer")
16class _DeferredShift(Exception):
17 """採色を保留し、同じ段階の位置合わせを集める。"""
20@dataclass
21class _TrackState:
22 bounds: tuple
23 dy: int = 0
24 dx: int = 0
25 back_y: int = 0
26 back_x: int = 0
29log = logging.getLogger(__name__)
32def _batch_capacity(rows, cols):
33 limit = max(4, min(32, 16 * 1024 * 1024 // (rows * cols * 8)))
34 return 1 << (limit.bit_length() - 1)
37def boundary_key(cell):
38 key = getattr(cell, "_alignment_key", None)
39 if key is None:
40 key = (cell.boundary.shape, blake2b(cell.boundary.tobytes(), digest_size=16).digest())
41 cell._alignment_key = key
42 return key
45class AlignmentCache:
46 """塗り色が変わっても、同じ境界・矩形のFFTを繰り返さない。"""
48 def __init__(self, *, gpu="auto", expected_cells=1):
49 self.results = {}
50 self.hits = 0
51 self.cpu_solves = 0
52 self.gpu_solves = 0
53 self.gpu_error = None
54 self.gpu = os.environ.get("PMA_AUTO_PAINT_GPU", gpu).lower()
55 self.expected_cells = expected_cells
56 self._cupy = None
57 self._gpu_failed = False
58 self._buckets = {}
59 self._executor = None
60 self.cpu_batches = 0
61 self.batched_solves = 0
62 self._requests = None
63 self.regions = None
64 self.region_error = None
65 self._samples = None
66 from .auto_motion import MotionMemory
67 self.motion = MotionMemory()
68 self.motion_certified = 0
69 self.shared_support_hits = 0
70 self._supports = {}
71 self._shared = OrderedDict()
72 self.support_checks = self.motion_trials = 0
73 self._support_enabled = self._certify_enabled = True
74 self._resident = None
75 from .auto_routing import PipelinePolicy
76 self.policy = PipelinePolicy()
78 def register_motion(self, target, source, flow, x0=0, y0=0):
79 self.motion.register((boundary_key(target), boundary_key(source)), flow, x0, y0)
81 def _learn_motion(self, pending):
82 if not self._support_enabled and not self._certify_enabled: 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true
83 return
84 for key, (target, guide, bounds, _margin) in pending.items():
85 if key in self.results: 85 ↛ 84line 85 didn't jump to line 84 because the condition on line 85 was always true
86 self.motion.learn((boundary_key(target), boundary_key(guide)), bounds, self.results[key])
87 result = self.results[key]
88 if self._support_enabled and getattr(result[0], 'gap', 0) > 2e-5 and min(
89 abs(result[1] - threshold) for threshold in (0.3, 0.4, 0.5, 0.6)) > 2e-5:
90 support = self._support_key(target, bounds, _margin)
91 if support is not None:
92 self._shared[(boundary_key(guide), support)] = result
93 while len(self._shared) > 16384: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true
94 self._shared.popitem(last=False)
95 if self.support_checks >= 256 and self.shared_support_hits * 100 < self.support_checks:
96 self._support_enabled = False
97 self._shared.clear()
98 self._supports.clear()
99 if self.motion_trials >= 128 and self.motion_certified * 16 < self.motion_trials:
100 self._certify_enabled = False
102 def _support_key(self, target, bounds, margin):
103 from . import auto_paint as ap
104 key = boundary_key(target), tuple(bounds), margin
105 if key in self._supports:
106 return self._supports[key]
107 height, width = target.values.shape
108 x0, y0, x1, y1 = bounds
109 tx0, ty0 = max(0, x0-ap.CONTEXT), max(0, y0-ap.CONTEXT)
110 tx1, ty1 = min(width, x1+ap.CONTEXT), min(height, y1+ap.CONTEXT)
111 wx0, wy0 = max(0, tx0-margin), max(0, ty0-margin)
112 wx1, wy1 = min(width, tx1+margin), min(height, ty1+margin)
113 step = max(1, int(np.ceil(max(wx1-wx0, wy1-wy0)/ap._MAX_WINDOW)))
114 limit = margin // step
115 if min((wx1-wx0+step-1)//step, (wy1-wy0+step-1)//step) <= 2*limit:
116 self._supports[key] = None
117 return None
118 sx, sy = tx0+(wx0-tx0)%step, ty0+(wy0-ty0)%step
119 ys, xs = np.nonzero(target.boundary[sy:ty1:step, sx:tx1:step])
120 points = np.ascontiguousarray((ys*step+sy)*width + xs*step+sx, dtype=np.int64)
121 result = (target.values.shape, step, limit,
122 blake2b(points.tobytes(), digest_size=16).digest())
123 self._supports[key] = result
124 return result
126 def enable_regions(self, pixels):
127 if (self.gpu in ("off", "0") or
128 (self.gpu == "auto" and (self.expected_cells < 4 or pixels < 256 * 256))):
129 return
130 if os.environ.get("PMA_AUTO_PAINT_GPU_REGIONS") == "0": 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true
131 return
132 if self._load_gpu(): 132 ↛ exitline 132 didn't return from function 'enable_regions' because the condition on line 132 was always true
133 from .auto_gpu_regions import RegionGPU
134 try:
135 self.regions = RegionGPU(self._cupy, mode=self.gpu)
136 except Exception as error:
137 self.region_error = str(error)
138 log.warning("GPU region initialization failed; continuing on CPU: %s", self.region_error)
140 def close(self):
141 if self._executor is not None:
142 self._executor.shutdown(wait=True, cancel_futures=True)
143 self._executor = None
144 if self.regions is not None: 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true
145 self.regions.close()
146 self._shared.clear()
147 if self._resident is not None:
148 # 結果をCPUのタプルへ戻してから、GPUの小さな結果バッファも解放する。
149 self.results.update({key: (tuple(result[0]), result[1])
150 for key, result in self.results.items()})
151 try:
152 self._resident.close()
153 except Exception as error:
154 self.gpu_error = str(error)
155 log.warning("GPU alignment cleanup failed: %s", self.gpu_error)
157 def sample_region(self, guide, filled, shift=(0, 0), *, ignore=0, require=0, by_label=False):
158 from . import native
159 backend = self.regions
160 if backend is not None and not backend.failed and getattr(guide, "_native_immutable", False): 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true
161 key = backend.key(guide, filled, shift, ignore, require, by_label)
162 cached = backend.cached(key)
163 if cached is not None:
164 return cached
165 request = (guide, filled, shift, ignore, require, by_label)
166 if self._samples is not None:
167 self._samples[key] = request
168 raise _DeferredShift
169 backend.sample_many({key: request})
170 cached = backend.cached(key)
171 if cached is not None:
172 return cached
173 return native.sample_region(guide, filled, shift, ignore=ignore, require=require, by_label=by_label)
175 def solve_many(self, requests):
176 """必要と確定した探索だけをまとめ、キャッシュへの書き込みは主スレッドで行う。"""
177 from . import auto_paint as ap
178 pending = {}
179 for target, guide, bounds, margin in requests:
180 key = self.key(target, guide, bounds, margin)
181 if key not in self.results and key not in pending:
182 pending[key] = (target, guide, bounds, margin)
183 if not pending:
184 return
185 for key, (target, guide, bounds, margin) in list(pending.items()):
186 if not self._support_enabled: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 break
188 self.support_checks += 1
189 support = self._support_key(target, bounds, margin)
190 shared = self._shared.get((boundary_key(guide), support)) if support is not None else None
191 if shared is not None:
192 self.results[key] = shared
193 self.shared_support_hits += 1
194 del pending[key]
195 if not pending:
196 return
198 hints = {key: self.motion.predict((boundary_key(target), boundary_key(guide)), bounds) if self._certify_enabled else None
199 for key, (target, guide, bounds, _margin) in pending.items()}
200 self.motion_trials += sum(hint is not None for hint in hints.values())
202 use_gpu = (self.gpu not in ("off", "0") and not self._gpu_failed and len(pending) >= 4
203 and (self.gpu != "auto" or self.expected_cells >= 4))
204 if use_gpu and self._load_gpu(): 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true
205 problems = {}
206 groups = defaultdict(list)
207 for key, (target, guide, bounds, margin) in pending.items():
208 problem = ap._prepare_local_shift(target, guide, bounds, margin=margin, lazy=True)
209 if problem is None:
210 self.results[key] = ((0, 0), 0.0)
211 self.cpu_solves += 1
212 continue
213 problem['share_motion'] = self._support_enabled
214 from .auto_motion import certify
215 certified = certify(problem, hints[key])
216 if certified is not None:
217 self.results[key] = certified
218 self.motion_certified += 1
219 continue
220 problems[key] = problem
221 rows, cols = problem["size"]
222 if min(rows, cols) <= 2 * problem['limit']:
223 continue # 小窓の循環相関はCPUの従来寸法で厳密に保持する。
224 bucket = ((rows + 63) // 64 * 64, (cols + 63) // 64 * 64, problem["limit"])
225 groups[bucket].append((key, problem))
226 for bucket, items in groups.items():
227 if len(items) < 4:
228 continue
229 capacity = _batch_capacity(*bucket[:2])
230 for start in range(0, len(items), capacity):
231 self._batch(bucket, items[start:start + capacity])
232 remaining = {key: problem for key, problem in problems.items() if key not in self.results}
233 if remaining:
234 if self._executor is None:
235 self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="alignment")
236 results = self._executor.map(ap._solve_local_shift, remaining.values())
237 for key, result in zip(remaining, results):
238 self.results[key] = result
239 self.cpu_solves += 1
240 self.cpu_batches += 1
241 self.batched_solves += len(pending)
242 self._learn_motion(pending)
243 return
245 def solve(item):
246 key, request = item
247 target, guide, bounds, margin = request
248 problem = ap._prepare_local_shift(target, guide, bounds, margin=margin)
249 if problem is None: 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true
250 return ((0, 0), 0.0), False
251 problem['share_motion'] = self._support_enabled
252 if hints[key] is not None:
253 from .auto_motion import certify
254 certified = certify(problem, hints[key])
255 if certified is not None:
256 return certified, True
257 return ap._solve_local_shift(problem), False
259 if len(pending) < 4:
260 results = map(solve, pending.items())
261 else:
262 if self._executor is None:
263 self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="alignment")
264 results = self._executor.map(solve, pending.items())
265 self.cpu_batches += 1
266 for key, (result, certified) in zip(pending, results):
267 self.results[key] = result
268 self.cpu_solves += not certified
269 self.motion_certified += certified
270 self.batched_solves += 1
271 self._learn_motion(pending)
273 def resolve(self, calls: list[Callable[[], _Answer]]) -> list[_Answer]:
274 """必要な探索を一括実行する。callsは再実行されるため確定結果の書き込みを含めない。"""
275 with self.regions.phase() if self.regions is not None else nullcontext():
276 answers: dict[int, _Answer] = {}
277 pending = list(enumerate(calls))
278 while pending:
279 unresolved = []
280 self._requests = {}
281 self._samples = {}
282 try:
283 for index, call in pending:
284 try:
285 answers[index] = call()
286 except _DeferredShift:
287 unresolved.append((index, call))
288 requests = list(self._requests.values())
289 samples = self._samples
290 finally:
291 self._requests = None
292 self._samples = None
293 self.solve_many(requests)
294 if samples and self.regions is not None: 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true
295 self.regions.sample_many(samples)
296 pending = unresolved
297 return [answers[index] for index in range(len(calls))]
299 def track_many(self, steps, bounds_list, shape):
300 """全区画を一段ずつ進め、復路の探索も往路が成功した区画だけでまとめる。"""
301 height, width = shape
302 states: list[_TrackState | None] = [_TrackState(bounds) for bounds in bounds_list]
303 for backward, pairs in ((False, steps), (True, list(reversed(steps)))):
304 for left, right in pairs:
305 requests, indices = [], []
306 for index, state in enumerate(states):
307 if state is None:
308 continue
309 bounds = state.bounds
310 clipped = (max(0, bounds[0]), max(0, bounds[1]),
311 min(width, bounds[2]), min(height, bounds[3]))
312 if clipped[2] <= clipped[0] or clipped[3] <= clipped[1]: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true
313 states[index] = None
314 continue
315 requests.append((right, left, clipped, 48) if backward else
316 (left, right, clipped, 48))
317 indices.append(index)
318 self.solve_many(requests)
319 for index, request in zip(indices, requests):
320 shift, score = self.shift(*request)
321 if score < 0.5:
322 states[index] = None
323 continue
324 state = states[index]
325 assert state is not None
326 if backward:
327 state.back_y += shift[0]
328 state.back_x += shift[1]
329 bounds = state.bounds
330 state.bounds = (bounds[0] + shift[1], bounds[1] + shift[0],
331 bounds[2] + shift[1], bounds[3] + shift[0])
332 else:
333 state.dy += shift[0]
334 state.dx += shift[1]
335 initial = bounds_list[index]
336 state.bounds = (initial[0] + state.dx, initial[1] + state.dy,
337 initial[2] + state.dx, initial[3] + state.dy)
338 return [(state.dy, state.dx, state.back_y, state.back_x) if state is not None else None
339 for state in states]
341 def key(self, target, guide, bounds, margin):
342 return boundary_key(target), boundary_key(guide), tuple(bounds), margin
344 def shift(self, target, guide, bounds, margin):
345 from . import auto_paint as ap
346 key = self.key(target, guide, bounds, margin)
347 if key not in self.results:
348 if self._requests is not None:
349 self._requests[key] = (target, guide, bounds, margin)
350 raise _DeferredShift
351 self.results[key] = ap._local_shift(target, guide, bounds, margin=margin)
352 self.cpu_solves += 1
353 else:
354 self.hits += 1
355 return self.results[key]
357 def prefetch(self, target, guides, bounds_list, margin):
358 """GPU用の問題をサイズごとにまとめる。CPU経路では必要なときだけ解く。"""
359 if self.gpu in ("off", "0") or self._gpu_failed: 359 ↛ 360line 359 didn't jump to line 360 because the condition on line 359 was never true
360 return
361 if self.gpu == "auto" and (self.expected_cells < 4 or target.values.size < 256 * 256):
362 return
363 if len(bounds_list) < 4: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 return
365 if not self._load_gpu(): 365 ↛ 366line 365 didn't jump to line 366 because the condition on line 365 was never true
366 return
367 from . import auto_paint as ap
368 pending = {}
369 for guide in guides:
370 for geometry in ([guide.initial, guide] if guide.initial is not None else [guide]):
371 for bounds in bounds_list:
372 key = self.key(target, geometry, bounds, margin)
373 if key in self.results or key in pending: 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true
374 continue
375 problem = ap._prepare_local_shift(target, geometry, bounds, margin=margin)
376 if problem is None: 376 ↛ 377line 376 didn't jump to line 377 because the condition on line 376 was never true
377 self.results[key] = ((0, 0), 0.0)
378 else:
379 pending[key] = problem
380 groups = defaultdict(list)
381 for key, problem in pending.items():
382 # 同じ演算を呼ぶたびにFFTプランを作らず、64画素単位でまとめる。
383 rows, cols = problem["size"]
384 bucket = ((rows + 63) // 64 * 64, (cols + 63) // 64 * 64, problem["limit"])
385 groups[bucket].append((key, problem))
386 for bucket, items in groups.items():
387 if len(items) < 4: 387 ↛ 388line 387 didn't jump to line 388 because the condition on line 387 was never true
388 continue # 小さなグループは転送せず、通常のCPU経路へ。
389 rows, cols, _limit = bucket
390 batch_size = _batch_capacity(rows, cols)
391 for start in range(0, len(items), batch_size):
392 batch = items[start:start + batch_size]
393 self._batch(bucket, batch)
395 def _load_gpu(self):
396 if self._cupy is not None: 396 ↛ 398line 396 didn't jump to line 398 because the condition on line 396 was always true
397 return True
398 try:
399 import importlib
400 cupy = importlib.import_module("cupy")
401 if cupy.cuda.runtime.getDeviceCount() == 0:
402 self._gpu_failed = True
403 return False
404 cache = cupy.fft.config.get_plan_cache()
405 cache.set_size(128)
406 _free, total = cupy.cuda.runtime.memGetInfo()
407 cache.set_memsize(min(512 * 1024 * 1024, total // 16))
408 self._cupy = cupy
409 return True
410 except Exception as error:
411 self.gpu_error = str(error)
412 self._gpu_failed = True
413 log.debug("Automatic paint uses CPU: %s", error)
414 return False
416 def _windows_gpu(self, bucket, batch):
417 cp = self._cupy
418 assert cp is not None
419 if self._resident is None:
420 from .auto_gpu_alignment import ResidentAlignment
421 self._resident = ResidentAlignment(cp)
422 return self._resident.windows(bucket, batch, _batch_capacity(*bucket[:2]))
424 def _compact_gpu(self, bucket, batch) -> tuple[np.ndarray, Any]:
425 from . import auto_paint as ap
426 windows = self._windows_gpu(bucket, batch)
427 if isinstance(windows, np.ndarray): 427 ↛ 437line 427 didn't jump to line 437 because the condition on line 427 was always true
428 # CPU配列を返す検証用バックエンドも同じ縮約規則を使う。
429 rows = []
430 for (_key, problem), window in zip(batch, windows):
431 shift, overlap = ap._finish_local_shift(problem, window)
432 scores = window / problem['total'] - ap.MOTION_PENALTY * problem['step'] * (
433 np.abs(problem['dys'])[:, None] + np.abs(problem['dxs'])[None, :])
434 best = np.partition(scores.ravel(), -2)[-2:]
435 rows.append((*shift, overlap, best[-1] - best[-2]))
436 return np.asarray(rows), None
437 assert self._resident is not None
438 return self._resident.reduce(batch, windows)
440 def _batch(self, bucket, batch):
441 from . import auto_paint as ap
442 if self._gpu_failed or self._buckets.get(bucket) is False: 442 ↛ 443line 442 didn't jump to line 443 because the condition on line 442 was never true
443 return
444 policy_key = ('alignment', bucket, len(batch).bit_length())
445 route = self.policy.choose(policy_key)
446 if bucket not in self._buckets and route is False:
447 return
448 try:
449 if bucket not in self._buckets and route is None:
450 # 最初のグループはCPUで解いて検証・計測する。初期化と転送を済ませて
451 # 再計測し、実際に速いサイズだけを以後GPUへ回す。
452 started = perf_counter()
453 if self._executor is None: 453 ↛ 455line 453 didn't jump to line 455 because the condition on line 453 was always true
454 self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="alignment")
455 cpu = list(self._executor.map(ap._solve_local_shift, [p for _key, p in batch]))
456 cpu_seconds = perf_counter() - started
457 started = perf_counter()
458 self._compact_gpu(bucket, batch)
459 cold_seconds = perf_counter() - started
460 started = perf_counter()
461 compact, _device = self._compact_gpu(bucket, batch)
462 warm_seconds = perf_counter() - started
463 gpu_seconds = warm_seconds + max(0.0, cold_seconds - warm_seconds) / max(1, self.expected_cells)
464 same_shifts = all(
465 (int(row[0]), int(row[1])) == expected[0]
466 for row, expected in zip(compact, cpu)
467 )
468 self.policy.record(policy_key, cpu_seconds, gpu_seconds, valid=same_shifts)
469 for (key, _problem), result in zip(batch, cpu):
470 self.results[key] = result
471 self.cpu_solves += len(batch)
472 return
473 compact, device = self._compact_gpu(bucket, batch)
474 for index, ((key, problem), row) in enumerate(zip(batch, compact)):
475 result = ((int(row[0]), int(row[1])), float(row[2]))
476 gap = float(row[3])
477 near_threshold = min(abs(result[1] - threshold)
478 for threshold in (ap.MIN_LINE_OVERLAP, ap.MIN_COLOR_LINE_OVERLAP,
479 0.3, 0.5, 0.6)) < 2e-5
480 if gap < 2e-5 or near_threshold: 480 ↛ 484line 480 didn't jump to line 484 because the condition on line 480 was always true
481 result = ap._solve_local_shift(problem)
482 self.cpu_solves += 1
483 else:
484 self.gpu_solves += 1
485 if device is not None:
486 from .auto_gpu_alignment import DeviceShift
487 result = (DeviceShift(*result[0], device[index, :2]), result[1])
488 result[0].gap = gap
489 self.results[key] = result
490 except Exception as error:
491 self._gpu_failed = True
492 self.gpu_error = str(error)
493 log.warning("Batched GPU alignment failed; continuing on CPU: %s", error)