Coverage for src / puzzletree / reconstruct / inspect.py: 57.55%

1from __future__ import annotations (empty)

2 

3from collections.abc import Sequence (empty)

4from dataclasses import dataclass (empty)

5from pathlib import Path (empty)

6from typing import Dict, List, Tuple (empty)

7 

8import numpy as np (empty)

9from PIL import Image, ImageDraw, ImageFont (empty)

10 

11from puzzletree.reconstruct.core import ( (empty)

12 AdjList, 

13 charged_path, 

14 chargeds, 

15 connected_components, 

16 edge_label_exists, 

17 global_edge_exists, 

18 ita_path, 

19) 

20from puzzletree.reconstruct.io import load_tiles_from_dir (empty)

21from puzzletree.reconstruct.pipeline import run_reconstruction (empty)

22from puzzletree.reconstruct.render import tile_rgba_images (empty)

23 

24SideName = str (empty)

25 

26FONT = ImageFont.load_default() (empty)

27TEXT_COLOR = (24, 24, 24) (empty)

28MUTED_TEXT_COLOR = (90, 90, 90) (empty)

29HIGHLIGHT_COLOR = (225, 76, 76) (empty)

30SECONDARY_HIGHLIGHT = (54, 125, 247) (empty)

31BORDER_COLOR = (200, 200, 200) (empty)

32BACKGROUND = (255, 255, 255) (empty)

33PANEL_BACKGROUND = (248, 248, 248) (empty)

34 

35 

36@dataclass(frozen=True) (empty)

37class CandidateInspection: (empty)

38 kind: str (empty)

39 source: int (empty)

40 target: int (empty)

41 weight: float (empty)

42 row_rank: int (empty)

43 col_rank: int (empty)

44 source_component: int (empty)

45 target_component: int (empty)

46 reasons: Tuple[str, ...] (empty)

47 

48 

49@dataclass(frozen=True) (empty)

50class InspectionDataset: (empty)

51 name: str (empty)

52 input_dir: Path (empty)

53 minset: float (empty)

54 tiles: List[np.ndarray] (empty)

55 adjs: AdjList (empty)

56 components: List[List[int]] (empty)

57 lr: np.ndarray (empty)

58 ud: np.ndarray (empty)

59 

60 

61def _measure_text(text: str) -> Tuple[int, int]: (empty)

62 dummy = Image.new("RGB", (1, 1), BACKGROUND) (empty)

63 draw = ImageDraw.Draw(dummy) (empty)

64 left, top, right, bottom = draw.textbbox((0, 0), text, font=FONT) (empty)

65 return int(right - left), int(bottom - top) (empty)

66 

67 

68def _text_block_height(lines: Sequence[str], line_gap: int = 4) -> int: (empty)

69 if not lines: 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true(empty)

70 return 0 

71 heights = [_measure_text(line)[1] for line in lines] (empty)

72 return sum(heights) + line_gap * (len(lines) - 1) (empty)

73 

74 

75def _draw_text_lines( (empty)

76 draw: ImageDraw.ImageDraw, 

77 xy: Tuple[int, int], 

78 lines: Sequence[str], 

79 fill: Tuple[int, int, int] = TEXT_COLOR, 

80 line_gap: int = 4, 

81) -> None: 

82 x, y = xy (empty)

83 for line in lines: (empty)

84 draw.text((x, y), line, fill=fill, font=FONT) (empty)

85 y += _measure_text(line)[1] + line_gap (empty)

86 

87 

88def _tile_box( (empty)

89 origin_x: int, 

90 origin_y: int, 

91 tile_w: int, 

92 tile_h: int, 

93 scale: int, 

94 grid_x: int, 

95 grid_y: int, 

96) -> Tuple[int, int, int, int]: 

97 x0 = origin_x + grid_x * tile_w * scale (empty)

98 y0 = origin_y + grid_y * tile_h * scale (empty)

99 return (x0, y0, x0 + tile_w * scale, y0 + tile_h * scale) (empty)

100 

101 

102def _draw_side_highlight( (empty)

103 draw: ImageDraw.ImageDraw, 

104 box: Tuple[int, int, int, int], 

105 side: SideName, 

106 color: Tuple[int, int, int], 

107) -> None: 

108 x0, y0, x1, y1 = box (empty)

109 width = max(4, (x1 - x0) // 24) (empty)

110 if side == "left": (empty)

111 draw.line((x0, y0, x0, y1), fill=color, width=width) (empty)

112 elif side == "right": 112 ↛ 114line 112 didn't jump to line 114 because the condition on line 112 was always true(empty)

113 draw.line((x1, y0, x1, y1), fill=color, width=width) (empty)

114 elif side == "top": 

115 draw.line((x0, y0, x1, y0), fill=color, width=width) 

116 elif side == "bottom": 

117 draw.line((x0, y1, x1, y1), fill=color, width=width) 

118 else: 

119 raise ValueError(f"Unknown side: {side}") 

120 

121 

122def _label_for_kind(kind: str) -> Tuple[SideName, SideName]: (empty)

123 if kind == "lr": 123 ↛ 125line 123 didn't jump to line 125 because the condition on line 123 was always true(empty)

124 return ("right", "left") (empty)

125 if kind == "ud": 

126 return ("bottom", "top") 

127 raise ValueError(f"Unknown candidate kind: {kind}") 

128 

129 

130def _component_coords(adjs: AdjList, component: Sequence[int]) -> Dict[int, Tuple[int, int]]: (empty)

131 root = component[0] (empty)

132 comp_set = set(component) (empty)

133 return {node: coord for node, coord in chargeds(adjs, root) if node in comp_set} (empty)

134 

135 

136def _render_component_panel( (empty)

137 tile_imgs: Sequence[Image.Image], 

138 adjs: AdjList, 

139 component: Sequence[int], 

140 highlight_node: int, 

141 highlight_side: SideName, 

142 title: str, 

143 subtitle: str, 

144 scale: int = 1, 

145 padding: int = 16, 

146) -> Image.Image: 

147 coords = _component_coords(adjs, component) (empty)

148 tile_w, tile_h = tile_imgs[0].size (empty)

149 xs = [coord[0] for coord in coords.values()] (empty)

150 ys = [coord[1] for coord in coords.values()] (empty)

151 minx, maxx = min(xs), max(xs) (empty)

152 miny, maxy = min(ys), max(ys) (empty)

153 

154 text_lines = [title, subtitle] (empty)

155 text_height = _text_block_height(text_lines) (empty)

156 body_w = (maxx - minx + 1) * tile_w * scale (empty)

157 body_h = (maxy - miny + 1) * tile_h * scale (empty)

158 width = body_w + padding * 2 (empty)

159 height = body_h + padding * 2 + text_height + 10 (empty)

160 image = Image.new("RGB", (width, height), PANEL_BACKGROUND) (empty)

161 draw = ImageDraw.Draw(image) (empty)

162 _draw_text_lines(draw, (padding, padding), text_lines) (empty)

163 

164 origin_x = padding (empty)

165 origin_y = padding + text_height + 10 (empty)

166 draw.rectangle((origin_x - 1, origin_y - 1, origin_x + body_w, origin_y + body_h), outline=BORDER_COLOR, width=1) (empty)

167 

168 for node, (gx, gy) in coords.items(): (empty)

169 norm_x = gx - minx (empty)

170 norm_y = maxy - gy (empty)

171 box = _tile_box(origin_x, origin_y, tile_w, tile_h, scale, norm_x, norm_y) (empty)

172 resized = tile_imgs[node].resize((tile_w * scale, tile_h * scale), Image.Resampling.NEAREST).convert("RGB") (empty)

173 image.paste(resized, box[:2]) (empty)

174 draw.rectangle(box, outline=(220, 220, 220), width=1) (empty)

175 if node == highlight_node: (empty)

176 draw.rectangle(box, outline=HIGHLIGHT_COLOR, width=3) (empty)

177 _draw_side_highlight(draw, box, highlight_side, HIGHLIGHT_COLOR) (empty)

178 label = f"{node}" (empty)

179 lw, lh = _measure_text(label) (empty)

180 label_box = (box[0] + 4, box[1] + 4, box[0] + 10 + lw, box[1] + 8 + lh) (empty)

181 draw.rectangle(label_box, fill=(255, 255, 255)) (empty)

182 draw.text((box[0] + 7, box[1] + 6), label, fill=TEXT_COLOR, font=FONT) (empty)

183 

184 return image (empty)

185 

186 

187def _render_pair_zoom_panel( (empty)

188 tile_imgs: Sequence[Image.Image], 

189 candidate: CandidateInspection, 

190 scale: int = 2, 

191 padding: int = 16, 

192) -> Image.Image: 

193 source_side, target_side = _label_for_kind(candidate.kind) (empty)

194 source_tile = ( (empty)

195 tile_imgs[candidate.source] 

196 .resize( 

197 (tile_imgs[candidate.source].width * scale, tile_imgs[candidate.source].height * scale), 

198 Image.Resampling.NEAREST, 

199 ) 

200 .convert("RGB") 

201 ) 

202 target_tile = ( (empty)

203 tile_imgs[candidate.target] 

204 .resize( 

205 (tile_imgs[candidate.target].width * scale, tile_imgs[candidate.target].height * scale), 

206 Image.Resampling.NEAREST, 

207 ) 

208 .convert("RGB") 

209 ) 

210 

211 if candidate.kind == "lr": 211 ↛ 216line 211 didn't jump to line 216 because the condition on line 211 was always true(empty)

212 gap = 18 (empty)

213 body_w = source_tile.width + target_tile.width + gap (empty)

214 body_h = max(source_tile.height, target_tile.height) (empty)

215 else: 

216 gap = 18 

217 body_w = max(source_tile.width, target_tile.width) 

218 body_h = source_tile.height + target_tile.height + gap 

219 

220 lines = [ (empty)

221 f"Candidate {candidate.kind}: {candidate.source} -> {candidate.target}", 

222 f"weight={candidate.weight:.6f} ranks=({candidate.row_rank},{candidate.col_rank})", 

223 f"reasons={','.join(candidate.reasons)}", 

224 ] 

225 text_height = _text_block_height(lines) (empty)

226 width = body_w + padding * 2 (empty)

227 height = body_h + padding * 2 + text_height + 10 (empty)

228 image = Image.new("RGB", (width, height), PANEL_BACKGROUND) (empty)

229 draw = ImageDraw.Draw(image) (empty)

230 _draw_text_lines(draw, (padding, padding), lines) (empty)

231 

232 origin_x = padding (empty)

233 origin_y = padding + text_height + 10 (empty)

234 if candidate.kind == "lr": 234 ↛ 243line 234 didn't jump to line 243 because the condition on line 234 was always true(empty)

235 left_box = (origin_x, origin_y, origin_x + source_tile.width, origin_y + source_tile.height) (empty)

236 right_box = ( (empty)

237 origin_x + source_tile.width + gap, 

238 origin_y, 

239 origin_x + source_tile.width + gap + target_tile.width, 

240 origin_y + target_tile.height, 

241 ) 

242 else: 

243 left_box = (origin_x, origin_y, origin_x + source_tile.width, origin_y + source_tile.height) 

244 right_box = ( 

245 origin_x, 

246 origin_y + source_tile.height + gap, 

247 origin_x + target_tile.width, 

248 origin_y + source_tile.height + gap + target_tile.height, 

249 ) 

250 

251 image.paste(source_tile, left_box[:2]) (empty)

252 image.paste(target_tile, right_box[:2]) (empty)

253 draw.rectangle(left_box, outline=SECONDARY_HIGHLIGHT, width=3) (empty)

254 draw.rectangle(right_box, outline=HIGHLIGHT_COLOR, width=3) (empty)

255 _draw_side_highlight(draw, left_box, source_side, SECONDARY_HIGHLIGHT) (empty)

256 _draw_side_highlight(draw, right_box, target_side, HIGHLIGHT_COLOR) (empty)

257 return image (empty)

258 

259 

260def _stack_horizontally( (empty)

261 images: Sequence[Image.Image], 

262 gap: int = 16, 

263 background: Tuple[int, int, int] = BACKGROUND, 

264) -> Image.Image: 

265 if not images: 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true(empty)

266 raise ValueError("No images to stack.") 

267 width = sum(image.width for image in images) + gap * (len(images) - 1) (empty)

268 height = max(image.height for image in images) (empty)

269 out = Image.new("RGB", (width, height), background) (empty)

270 x = 0 (empty)

271 for image in images: (empty)

272 y = (height - image.height) // 2 (empty)

273 out.paste(image, (x, y)) (empty)

274 x += image.width + gap (empty)

275 return out (empty)

276 

277 

278def _stack_vertically( (empty)

279 images: Sequence[Image.Image], 

280 gap: int = 18, 

281 background: Tuple[int, int, int] = BACKGROUND, 

282) -> Image.Image: 

283 if not images: 283 ↛ 284line 283 didn't jump to line 284 because the condition on line 283 was never true(empty)

284 raise ValueError("No images to stack.") 

285 width = max(image.width for image in images) (empty)

286 height = sum(image.height for image in images) + gap * (len(images) - 1) (empty)

287 out = Image.new("RGB", (width, height), background) (empty)

288 y = 0 (empty)

289 for image in images: (empty)

290 x = (width - image.width) // 2 (empty)

291 out.paste(image, (x, y)) (empty)

292 y += image.height + gap (empty)

293 return out (empty)

294 

295 

296def _resize_to_fit(image: Image.Image, max_width: int = 4200, max_height: int = 9000) -> Image.Image: (empty)

297 scale = min(max_width / image.width, max_height / image.height, 1.0) 

298 if scale == 1.0: 

299 return image 

300 resized = image.resize( 

301 (max(1, int(image.width * scale)), max(1, int(image.height * scale))), 

302 Image.Resampling.BICUBIC, 

303 ) 

304 return resized 

305 

306 

307def load_inspection_dataset(input_dir: Path, minset: float = 0.1, r: float = 12.0) -> InspectionDataset: (empty)

308 tiles = load_tiles_from_dir(input_dir) 

309 result = run_reconstruction(tiles, r=r, minset=minset) 

310 from puzzletree.reconstruct.core import build_weight_matrices 

311 

312 lr, ud = build_weight_matrices(tiles, r=r) 

313 components = connected_components(result.adjs) 

314 return InspectionDataset( 

315 name=input_dir.stem.replace("_tiles", ""), 

316 input_dir=input_dir, 

317 minset=minset, 

318 tiles=tiles, 

319 adjs=result.adjs, 

320 components=components, 

321 lr=lr, 

322 ud=ud, 

323 ) 

324 

325 

326def inspect_candidate(dataset: InspectionDataset, kind: str, source: int, target: int) -> CandidateInspection: (empty)

327 components = dataset.components 

328 source_component = next(index for index, component in enumerate(components) if source in component) 

329 target_component = next(index for index, component in enumerate(components) if target in component) 

330 if source_component == target_component: 

331 raise ValueError("Candidate is within one component; expected cross-tree candidate.") 

332 

333 matrix = dataset.lr if kind == "lr" else dataset.ud 

334 weight = float(matrix[source, target]) 

335 row_rank = int(np.argsort(matrix[source]).tolist().index(target)) + 1 

336 col_rank = int(np.argsort(matrix[:, target]).tolist().index(source)) + 1 

337 

338 directions = (-1, 0) if kind == "lr" else (0, 1) 

339 labels = [1, 2] if kind == "lr" else [3, 4] 

340 location1, location2 = source, target 

341 if location1 < location2: 

342 location1, location2 = location2, location1 

343 labels = [labels[1], labels[0]] 

344 directions = (-directions[0], -directions[1]) 

345 

346 reasons: List[str] = [] 

347 if weight > dataset.minset: 

348 reasons.append("above_minset") 

349 if edge_label_exists(dataset.adjs[location1], labels[-1]): 

350 reasons.append("side_used_location1") 

351 if global_edge_exists(dataset.adjs, labels[-1], location2): 

352 reasons.append("global_side_used_location2") 

353 if global_edge_exists(dataset.adjs, labels[0], location1): 

354 reasons.append("global_side_used_location1") 

355 if edge_label_exists(dataset.adjs[location2], labels[0]): 

356 reasons.append("side_used_location2") 

357 if ita_path(dataset.adjs, location1, location2): 

358 reasons.append("same_component_path") 

359 

360 charges_l2 = [coord for _, coord in chargeds(dataset.adjs, location2)] 

361 overlap_l1 = any( 

362 charged_path(dataset.adjs, location1, (coord[0] - directions[0], coord[1] - directions[1])) 

363 for coord in charges_l2 

364 ) 

365 charges_l1 = [coord for _, coord in chargeds(dataset.adjs, location1)] 

366 overlap_l2 = any( 

367 charged_path(dataset.adjs, location2, (coord[0] + directions[0], coord[1] + directions[1])) 

368 for coord in charges_l1 

369 ) 

370 if overlap_l1 or overlap_l2: 

371 reasons.append("overlap") 

372 

373 if not reasons: 

374 reasons.append("would_accept") 

375 

376 return CandidateInspection( 

377 kind=kind, 

378 source=source, 

379 target=target, 

380 weight=weight, 

381 row_rank=row_rank, 

382 col_rank=col_rank, 

383 source_component=source_component, 

384 target_component=target_component, 

385 reasons=tuple(reasons), 

386 ) 

387 

388 

389def find_low_score_rejected_candidates(dataset: InspectionDataset, limit: int = 6) -> List[CandidateInspection]: (empty)

390 candidates: List[CandidateInspection] = [] 

391 seen: set[Tuple[Tuple[int, int], str]] = set() 

392 node_to_component = {node: index for index, component in enumerate(dataset.components) for node in component} 

393 

394 weighted_candidates: List[Tuple[str, float, int, int]] = [] 

395 for source in range(len(dataset.tiles)): 

396 for target in range(len(dataset.tiles)): 

397 if source == target: 

398 continue 

399 if node_to_component[source] == node_to_component[target]: 

400 continue 

401 weighted_candidates.append(("lr", float(dataset.lr[source, target]), source, target)) 

402 weighted_candidates.append(("ud", float(dataset.ud[source, target]), source, target)) 

403 weighted_candidates.sort(key=lambda item: item[1]) 

404 

405 for kind, _, source, target in weighted_candidates: 

406 inspection = inspect_candidate(dataset, kind, source, target) 

407 if inspection.reasons == ("would_accept",): 

408 continue 

409 pair = ( 

410 min(inspection.source_component, inspection.target_component), 

411 max(inspection.source_component, inspection.target_component), 

412 ) 

413 key = (pair, kind) 

414 if key in seen: 

415 continue 

416 seen.add(key) 

417 candidates.append(inspection) 

418 if len(candidates) >= limit: 

419 break 

420 return candidates 

421 

422 

423def render_candidate_panel(dataset: InspectionDataset, candidate: CandidateInspection) -> Image.Image: (empty)

424 tile_imgs = tile_rgba_images(dataset.tiles) (empty)

425 source_side, target_side = _label_for_kind(candidate.kind) (empty)

426 left_component = dataset.components[candidate.source_component] (empty)

427 right_component = dataset.components[candidate.target_component] (empty)

428 

429 left_panel = _render_component_panel( (empty)

430 tile_imgs, 

431 dataset.adjs, 

432 left_component, 

433 highlight_node=candidate.source, 

434 highlight_side=source_side, 

435 title=f"Tree {candidate.source_component} size={len(left_component)}", 

436 subtitle=f"tiles={sorted(left_component)}", 

437 ) 

438 center_panel = _render_pair_zoom_panel(tile_imgs, candidate) (empty)

439 right_panel = _render_component_panel( (empty)

440 tile_imgs, 

441 dataset.adjs, 

442 right_component, 

443 highlight_node=candidate.target, 

444 highlight_side=target_side, 

445 title=f"Tree {candidate.target_component} size={len(right_component)}", 

446 subtitle=f"tiles={sorted(right_component)}", 

447 ) 

448 return _stack_horizontally([left_panel, center_panel, right_panel], gap=20) (empty)

449 

450 

451def render_candidate_contact_sheet( (empty)

452 dataset: InspectionDataset, 

453 candidates: Sequence[CandidateInspection], 

454) -> Image.Image: 

455 title_lines = [ (empty)

456 f"Dataset: {dataset.name}", 

457 f"minset={dataset.minset} components={len(dataset.components)} sizes={sorted((len(c) for c in dataset.components), reverse=True)}", 

458 "Each row shows left tree, candidate edge zoom, and right tree. Red/blue lines mark the compared sides.", 

459 ] 

460 header_h = _text_block_height(title_lines) + 32 (empty)

461 rows = [render_candidate_panel(dataset, candidate) for candidate in candidates] (empty)

462 body = _stack_vertically(rows, gap=20) (empty)

463 width = body.width + 32 (empty)

464 height = body.height + header_h + 16 (empty)

465 out = Image.new("RGB", (width, height), BACKGROUND) (empty)

466 draw = ImageDraw.Draw(out) (empty)

467 _draw_text_lines(draw, (16, 16), title_lines) (empty)

468 out.paste(body, (16, header_h)) (empty)

469 return out (empty)

470 

471 

472def save_candidate_report( (empty)

473 input_dir: Path, 

474 output_path: Path, 

475 minset: float = 0.1, 

476 r: float = 12.0, 

477 limit: int = 6, 

478) -> List[CandidateInspection]: 

479 dataset = load_inspection_dataset(input_dir=input_dir, minset=minset, r=r) 

480 candidates = find_low_score_rejected_candidates(dataset, limit=limit) 

481 image = render_candidate_contact_sheet(dataset, candidates) 

482 image = _resize_to_fit(image) 

483 output_path.parent.mkdir(parents=True, exist_ok=True) 

484 image.save(output_path) 

485 return candidates