Coverage for src / puzzletree / reconstruct / pipeline.py: 100.00%

1from __future__ import annotations (empty)

2 

3from dataclasses import dataclass (empty)

4from pathlib import Path (empty)

5from typing import Callable, List, cast (empty)

6 

7import numpy as np (empty)

8from PIL import Image (empty)

9 

10from puzzletree.reconstruct.core import AdjList, Coord, build_weight_matrices, msgt, reconstruct_layout (empty)

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

12from puzzletree.reconstruct.render import render_reconstruction, save_tree_build_animation (empty)

13 

14ProgressCallback = Callable[[str], None] (empty)

15 

16 

17@dataclass (empty)

18class ReconstructionRun: (empty)

19 adjs: AdjList (empty)

20 placements: dict[int, Coord] (empty)

21 output: Image.Image (empty)

22 

23 

24@dataclass (empty)

25class ReconstructionRunWithHistory(ReconstructionRun): (empty)

26 history: List[AdjList] (empty)

27 

28 

29@dataclass (empty)

30class ReconstructOptions: (empty)

31 input_dir: Path (empty)

32 output: Path = Path("reconstructed.png") (empty)

33 r: float = 12.0 (empty)

34 minset: float = 0.1 (empty)

35 animation: Path | None = None (empty)

36 animation_seed: int = 0 (empty)

37 animation_size: int = 1024 (empty)

38 animation_max_angle: float = 35.0 (empty)

39 animation_duration_ms: int = 1000 (empty)

40 animation_frames_dir: Path | None = None (empty)

41 

42 

43def _notify_progress(progress_callback: ProgressCallback | None, stage: str) -> None: (empty)

44 if progress_callback is not None: (empty)

45 progress_callback(stage) (empty)

46 

47 

48def run_reconstruction( (empty)

49 tiles: List[np.ndarray], 

50 r: float, 

51 minset: float, 

52 progress_callback: ProgressCallback | None = None, 

53) -> ReconstructionRun: 

54 h, w = tiles[0].shape[:2] (empty)

55 _notify_progress(progress_callback, "Computing edge weights") (empty)

56 lr, ud = build_weight_matrices(tiles, r=r) (empty)

57 _notify_progress(progress_callback, "Assembling reconstruction tree") (empty)

58 adjs = cast("AdjList", msgt(lr, ud, minset=minset, lr_side_size=w, ud_side_size=h)) (empty)

59 _notify_progress(progress_callback, "Rendering reconstructed image") (empty)

60 placements = reconstruct_layout(adjs) (empty)

61 output = render_reconstruction(tiles, placements) (empty)

62 return ReconstructionRun(adjs=adjs, placements=placements, output=output) (empty)

63 

64 

65def run_reconstruction_with_history( (empty)

66 tiles: List[np.ndarray], 

67 r: float, 

68 minset: float, 

69 progress_callback: ProgressCallback | None = None, 

70) -> ReconstructionRunWithHistory: 

71 h, w = tiles[0].shape[:2] (empty)

72 _notify_progress(progress_callback, "Computing edge weights") (empty)

73 lr, ud = build_weight_matrices(tiles, r=r) (empty)

74 _notify_progress(progress_callback, "Assembling reconstruction tree") (empty)

75 adjs, history = cast( (empty)

76 "tuple[AdjList, List[AdjList]]", 

77 msgt(lr, ud, minset=minset, lr_side_size=w, ud_side_size=h, record_history=True), 

78 ) 

79 _notify_progress(progress_callback, "Rendering reconstructed image") (empty)

80 placements = reconstruct_layout(adjs) (empty)

81 output = render_reconstruction(tiles, placements) (empty)

82 return ReconstructionRunWithHistory(adjs=adjs, placements=placements, output=output, history=history) (empty)

83 

84 

85def run_from_options( (empty)

86 options: ReconstructOptions, 

87 progress_callback: ProgressCallback | None = None, 

88) -> ReconstructionRun | ReconstructionRunWithHistory: 

89 _notify_progress(progress_callback, "Loading tiles") (empty)

90 tiles = load_tiles_from_dir(options.input_dir) (empty)

91 

92 if options.animation is None: (empty)

93 result = run_reconstruction(tiles, r=options.r, minset=options.minset, progress_callback=progress_callback) (empty)

94 else: 

95 result = run_reconstruction_with_history( (empty)

96 tiles, 

97 r=options.r, 

98 minset=options.minset, 

99 progress_callback=progress_callback, 

100 ) 

101 

102 _notify_progress(progress_callback, "Saving reconstructed image") (empty)

103 result.output.save(options.output) (empty)

104 

105 if options.animation is not None and isinstance(result, ReconstructionRunWithHistory): (empty)

106 _notify_progress(progress_callback, "Rendering animation") (empty)

107 save_tree_build_animation( (empty)

108 tiles=tiles, 

109 history=result.history, 

110 output_path=options.animation, 

111 seed=options.animation_seed, 

112 frame_size=options.animation_size, 

113 max_angle=options.animation_max_angle, 

114 duration_ms=options.animation_duration_ms, 

115 frames_dir=options.animation_frames_dir, 

116 ) 

117 

118 return result (empty)