Coverage for paintmaskanimator/native.py: 97%

1"""Required Rust backend for timeline state and image computations. 

2 

3Missing or outdated extensions fail at startup with a rebuild instruction. 

4NumPy remains in use for array conversion and algorithms not yet moved to Rust. 

5""" 

6from __future__ import annotations 

7 

8import numpy as np 

9 

10from .logging_setup import get_logger 

11 

12log = get_logger(__name__) 

13 

14REQUIRED_API = ( 

15 "TimelineCell", "resolve_timeline", "ClipAnimationReader", "GuideSampler", 

16 "resolve_exposures", "tint_onion", "unique_colors", "filter_palette", 

17 "line_flow_scores", "connect_runs", "label_runs", "refine_face", "thin", 

18 "render_quad_labels", "render_mesh_labels", 

19 "surface_guided_palette_labels", "map_known_colors", "region_neighbors", 

20 "flood_fill", "dilate", "erode", "label_regions", 

21) 

22REBUILD_MESSAGE = ( 

23 "PaintMaskAnimator requires its Rust extension. " 

24 "Install or rebuild it from the project directory: python -m pip install ./native" 

25) 

26 

27 

28def validate_backend(module): 

29 missing = [name for name in REQUIRED_API if not callable(getattr(module, name, None))] 

30 if missing: 

31 raise ImportError(f"{REBUILD_MESSAGE} (missing API: {', '.join(missing)})") 

32 

33 

34try: 

35 import pma_native as _native 

36except ImportError as error: 

37 raise ImportError(REBUILD_MESSAGE) from error 

38validate_backend(_native) 

39 

40 

41def available() -> bool: 

42 """Compatibility probe: the required backend was validated on import.""" 

43 return True 

44 

45 

46def clip_animation_reader(path: str, max_dimension: int, max_raster_bytes: int): 

47 """CLIP parsing has one implementation, supplied by clipfile-rs.""" 

48 return _native.ClipAnimationReader(path, max_dimension, max_raster_bytes) 

49 

50 

51def resolve_exposures(states, exposures): 

52 count, layers = states.shape 

53 data = _native.resolve_exposures( 

54 states.tobytes(), exposures.tobytes(), count, layers, 

55 ) 

56 return np.frombuffer(data, dtype=np.int64).reshape(count, layers) 

57 

58 

59def resolve_timeline(frames): 

60 """Read Rust-owned cell metadata directly, without packing NumPy inputs.""" 

61 rows = [[layer.__dict__.get("_timeline_metadata") for layer in frame.layers] 

62 for frame in frames] 

63 layers = max(map(len, rows), default=0) 

64 data = _native.resolve_timeline(rows) 

65 return np.frombuffer(data, dtype=np.int64).reshape(len(rows), layers) 

66 

67 

68def tint_onion(rgba, color, color_enabled, selected): 

69 data = _native.tint_onion( 

70 np.ascontiguousarray(rgba).tobytes(), tuple(color), bool(color_enabled), 

71 None if selected is None else [(int(rgb[0]), int(rgb[1]), int(rgb[2])) for rgb in selected], 

72 ) 

73 return np.frombuffer(data, dtype=np.uint8).reshape(rgba.shape) 

74 

75 

76def unique_colors(pixels, limit): 

77 """Return the smallest sorted packed RGB colors, using the required Rust backend.""" 

78 return _native.unique_colors( 

79 np.ascontiguousarray(pixels, dtype=np.uint32).tobytes(), int(limit) 

80 ) 

81 

82 

83def filter_palette(rgba, visible, legacy, remap): 

84 """Filter original RGB colors and remap visible pixels in one Rust pass.""" 

85 def rgb_tuple(rgb): 

86 return (int(rgb[0]), int(rgb[1]), int(rgb[2])) 

87 

88 data = _native.filter_palette( 

89 np.ascontiguousarray(rgba, dtype=np.uint8).tobytes(), 

90 None if visible is None else [rgb_tuple(rgb) for rgb in visible], 

91 None if legacy is None else rgb_tuple(legacy), 

92 [(rgb_tuple(child), rgb_tuple(parent)) for child, parent in remap.items()], 

93 ) 

94 return np.frombuffer(data, dtype=np.uint8).reshape(rgba.shape) 

95 

96 

97def line_flow_scores(target, guide, block, base_y, base_x, reach): 

98 """位置合わせの一致数をRustで計算する。""" 

99 height, width = target.shape 

100 result = _native.line_flow_scores( 

101 _as_bytes(target), _as_bytes(guide), width, height, int(block), 

102 np.ascontiguousarray(base_y, dtype=np.int64).tobytes(), 

103 np.ascontiguousarray(base_x, dtype=np.int64).tobytes(), int(reach), 

104 ) 

105 side = 2 * reach + 1 

106 return np.frombuffer(result, dtype=np.uint32).reshape( 

107 side * side, height // block, width // block, 

108 ).astype(np.float64) 

109 

110 

111def _as_bytes(mask): 

112 return np.ascontiguousarray(mask, dtype=bool).tobytes() 

113 

114 

115def connect_runs(starts, ends, kinds, rows): 

116 """同色の上下ランをRustで結合する。""" 

117 data, count = _native.connect_runs(*( 

118 np.ascontiguousarray(values, dtype=np.uint64).tobytes() 

119 for values in (starts, ends, kinds, rows) 

120 )) 

121 return np.frombuffer(data, dtype=np.int32).copy(), count 

122 

123 

124def label_runs(values): 

125 """連結領域の走査・結合・番号画像の作成をRustでまとめて行う。""" 

126 # 大画像は転送・番号画像の複製が支配的になるため、NumPyの走査と 

127 # connect_runsのRust結合を併用する方が速い（実データで確認）。 

128 if (values.dtype.kind not in "biu" or not values.dtype.isnative 

129 or values.size > 65536): 

130 return None 

131 height, width = values.shape 

132 labels, count, ys, xs, ends, runs = _native.label_runs( 

133 np.ascontiguousarray(values).tobytes(), width, height, values.dtype.itemsize, 

134 ) 

135 return (np.frombuffer(labels, dtype=np.int32).reshape(height, width).copy(), count, 

136 (np.frombuffer(ys, dtype=np.int64).copy(), np.frombuffer(xs, dtype=np.int64).copy(), 

137 np.frombuffer(ends, dtype=np.int64).copy(), np.frombuffer(runs, dtype=np.int32).copy())) 

138 

139 

140def refine_face(guide, filled, shift, bounds, radius, thresholds): 

141 """必要な小矩形だけを渡し、ずれの探索・採色・面積比確認を一括実行する。""" 

142 ax0, ay0, ax1, ay1 = bounds 

143 ys, xs = np.nonzero(filled.local) 

144 points = np.column_stack((ys + filled.bounds[1] + shift[0] - ay0, 

145 xs + filled.bounds[0] + shift[1] - ax0)) 

146 roi = np.s_[ay0:ay1, ax0:ax1] 

147 return _native.refine_face( 

148 np.ascontiguousarray(guide.values[roi], dtype=np.uint64).tobytes(), 

149 np.ascontiguousarray(guide.regions.labels[roi], dtype=np.uint32).tobytes(), 

150 _as_bytes(~guide.initial.boundary[roi]), 

151 _as_bytes(guide.open[roi] | guide.undecided[roi]), 

152 np.ascontiguousarray(guide.regions.sizes, dtype=np.uint64).tobytes(), 

153 np.ascontiguousarray(points, dtype=np.int64).tobytes(), 

154 ax1 - ax0, ay1 - ay0, radius, thresholds, 

155 ) 

156 

157 

158def sample_region(guide, filled, shift=(0, 0), *, ignore=0, require=0, by_label=False): 

159 """不変の見本をラン圧縮して保持し、面の重なりと色をRustで数える。""" 

160 immutable = getattr(guide, "_native_immutable", False) 

161 sampler = getattr(guide, "_native_sampler", None) if immutable else None 

162 if sampler is None: 

163 initial = guide.initial if guide.initial is not None else guide 

164 flags = guide.boundary.astype(np.uint8) 

165 flags |= initial.boundary.astype(np.uint8) * 2 

166 flags |= (guide.open | guide.undecided).astype(np.uint8) * 4 

167 confirmed = getattr(guide, "confirmed", None) 

168 flags |= 8 if confirmed is None else confirmed.astype(np.uint8) * 8 

169 sampler = _native.GuideSampler( 

170 np.ascontiguousarray(guide.values, dtype=np.uint64).tobytes(), 

171 np.ascontiguousarray(guide.regions.labels, dtype=np.uint32).tobytes(), 

172 flags.tobytes(), 

173 np.ascontiguousarray(guide.regions.sizes, dtype=np.uint64).tobytes(), 

174 guide.values.shape[1], guide.values.shape[0], 

175 ) 

176 if immutable: 

177 guide._native_sampler = sampler 

178 runs = getattr(filled, "_native_runs", None) 

179 if runs is None: 

180 padded = np.pad(filled.local, ((0, 0), (1, 1))) 

181 ys, xs = np.nonzero(padded[:, 1:] != padded[:, :-1]) 

182 points = np.column_stack((ys[::2] + filled.bounds[1], xs[::2] + filled.bounds[0], 

183 xs[1::2] + filled.bounds[0])) 

184 runs = np.ascontiguousarray(points, dtype=np.int64).tobytes() 

185 filled._native_runs = runs 

186 return sampler.sample(runs, int(shift[0]), int(shift[1]), ignore, require, by_label) 

187 

188 

189def thin(mask, iterations): 

190 """線端検出用の細線化をRustで計算する。""" 

191 height, width = mask.shape 

192 return _as_mask(_native.thin(_as_bytes(mask), width, height, int(iterations)), height, width) 

193 

194 

195def _as_mask(data, height, width): 

196 # 返ってきた bytes は読み取り専用なので、呼び出し側が書き換えられるよう複製する。 

197 return np.frombuffer(data, dtype=bool).reshape(height, width).copy() 

198 

199 

200def flood_fill(passable, seeds): 

201 """4近傍で ``seeds`` からつながる ``passable`` の領域。""" 

202 height, width = passable.shape 

203 seeds = [(int(x), int(y)) for x, y in seeds] 

204 return _as_mask( 

205 _native.flood_fill(_as_bytes(passable), width, height, seeds), height, width 

206 ) 

207 

208 

209def dilate(mask, iterations): 

210 """8近傍で ``iterations`` 回膨張させる。""" 

211 height, width = mask.shape 

212 return _as_mask( 

213 _native.dilate(_as_bytes(mask), width, height, int(iterations)), height, width 

214 ) 

215 

216 

217def erode(mask, iterations): 

218 """8近傍で ``iterations`` 回収縮させる（範囲外は偽）。""" 

219 height, width = mask.shape 

220 return _as_mask( 

221 _native.erode(_as_bytes(mask), width, height, int(iterations)), height, width 

222 ) 

223 

224 

225#: :func:`label_regions` で番号を付けない画素の値。 

226INVALID_CODE = 0xFFFFFFFF 

227 

228 

229def label_regions(codes): 

230 """同じ値で4近傍につながった画素に1からの番号を付ける。 

231 

232 :param codes: uint32 の2次元配列。:data:`INVALID_CODE` の画素は番号 0。 

233 :returns: ``(番号の int32 配列, 領域の数)``。 

234 """ 

235 height, width = codes.shape 

236 data, count = _native.label_regions( 

237 np.ascontiguousarray(codes, dtype=np.uint32).tobytes(), width, height 

238 ) 

239 labels = np.frombuffer(data, dtype=np.uint32).reshape(height, width) 

240 return labels.astype(np.int32), int(count) 

241 

242 

243def render_quad_labels(label_image, matrix, out_width, out_height, subsamples, coverage_labels): 

244 """平面四角形の逆写像で番号の画像を変形する。 

245 

246 :param label_image: uint16 の2次元配列。0 は「何もない」。 

247 :param matrix: 出力座標から元座標への 3x3（逆写像）。 

248 :param coverage_labels: 被覆（0-255）も返す番号の並び。 

249 :returns: ``(uint16 の番号, 番号 -> uint8 の被覆)``。 

250 """ 

251 height, width = label_image.shape 

252 labels = np.ascontiguousarray(label_image, dtype=np.uint16) 

253 coverage_labels = [int(label) for label in coverage_labels] 

254 data, coverage = _native.render_quad_labels( 

255 labels.tobytes(), width, height, 

256 [float(value) for value in np.asarray(matrix, dtype=np.float64).reshape(9)], 

257 int(out_width), int(out_height), int(subsamples), coverage_labels, 

258 ) 

259 result = np.frombuffer(data, dtype=np.uint16).reshape(out_height, out_width) 

260 return result.copy(), { 

261 label: np.frombuffer(plane, dtype=np.uint8).reshape(out_height, out_width).copy() 

262 for label, plane in zip(coverage_labels, coverage) 

263 } 

264 

265 

266def render_mesh_labels( 

267 label_image, points, dense_cols, dense_rows, out_width, out_height, 

268 subsamples, band_rows, coverage_labels, 

269): 

270 """曲がった格子（メッシュ変形）で番号の画像を変形する。 

271 

272 :param points: 出力座標の格子点 ``(x, y)`` を行優先で並べたもの。 

273 :param band_rows: 一度に解く行数。使う記憶領域の上限になる。 

274 :returns: ``(uint16 の番号, 番号 -> uint8 の被覆)``。 

275 """ 

276 height, width = label_image.shape 

277 labels = np.ascontiguousarray(label_image, dtype=np.uint16) 

278 coverage_labels = [int(label) for label in coverage_labels] 

279 grid = np.asarray(points, dtype=np.float64).reshape(-1, 2) 

280 data, coverage = _native.render_mesh_labels( 

281 labels.tobytes(), width, height, 

282 [(float(x), float(y)) for x, y in grid], 

283 int(dense_cols), int(dense_rows), int(out_width), int(out_height), 

284 int(subsamples), int(band_rows), coverage_labels, 

285 ) 

286 result = np.frombuffer(data, dtype=np.uint16).reshape(out_height, out_width) 

287 return result.copy(), { 

288 label: np.frombuffer(plane, dtype=np.uint8).reshape(out_height, out_width).copy() 

289 for label, plane in zip(coverage_labels, coverage) 

290 } 

291 

292 

293def surface_guided_palette_labels(rgb, visible, palette, spatial_refinement=True): 

294 """画素ごとに割り当てるパレットの番号を決める（色面を種に境界を二分する）。 

295 

296 :param rgb: ``(高さ, 幅, 3)`` の uint8 配列。 

297 :param visible: ``(高さ, 幅)`` の真偽配列。偽の画素は -1 のまま。 

298 :param palette: ``(色数, 3)`` の uint8 配列。 

299 :returns: ``(高さ, 幅)`` の int16 配列。未確定は -1。 

300 """ 

301 height, width = visible.shape 

302 pixels = np.ascontiguousarray(rgb, dtype=np.uint8).reshape(height, width, 3) 

303 colors = np.ascontiguousarray(palette, dtype=np.uint8).reshape(-1, 3) 

304 data = _native.surface_guided_palette_labels( 

305 pixels.tobytes(), width, height, 

306 np.ascontiguousarray(visible, dtype=bool).tobytes(), colors.tobytes(), 

307 bool(spatial_refinement), 

308 ) 

309 return np.frombuffer(data, dtype=np.int16).reshape(height, width).copy() 

310 

311 

312def map_known_colors(pixels, palette, color_mapping, line_max): 

313 """隣接色分けの仮の色を、割り当て済みの対応表どおりに塗る。 

314 

315 :param pixels: ``(高さ, 幅, 4)`` の RGBA8888 配列。 

316 :param color_mapping: 元の RGB をまとめた値から仮の色の番号への辞書。 

317 :returns: 塗った ``(高さ, 幅, 4)`` の配列。対応のない色があれば ``None``。 

318 """ 

319 height, width = pixels.shape[:2] 

320 codes = sorted(int(code) for code in color_mapping) 

321 data = _native.map_known_colors( 

322 np.ascontiguousarray(pixels, dtype=np.uint8).tobytes(), 

323 codes, [int(color_mapping[code]) for code in codes], 

324 [(int(color[0]), int(color[1]), int(color[2])) for color in palette], 

325 int(line_max), 

326 ) 

327 if data is None: 

328 return None 

329 # 読み取り専用のまま返す。呼び出し側が自分の配列へ書き写すので複製は要らない。 

330 return np.frombuffer(data, dtype=np.uint8).reshape(height, width, 4) 

331 

332 

333def region_neighbors(labels, gap): 

334 """``gap`` 画素以内で向かい合う領域の組 ``(a, b)``（``a < b``）の配列。""" 

335 height, width = labels.shape 

336 pairs = _native.region_neighbors( 

337 np.ascontiguousarray(labels, dtype=np.int32).tobytes(), width, height, int(gap) 

338 ) 

339 return np.asarray(pairs, dtype=np.int64).reshape(-1, 2)