Coverage for src / puzzletree / reconstruct / render.py: 87.80%

1from __future__ import annotations (empty)

2 

3import math (empty)

4import random (empty)

5from collections.abc import Sequence (empty)

6from pathlib import Path (empty)

7from typing import List, Tuple (empty)

8 

9import numpy as np (empty)

10from PIL import Image (empty)

11 

12from puzzletree.reconstruct.core import AdjList, chargeds, connected_components, reconstruct_layout (empty)

13 

14 

15def render_reconstruction(tiles: Sequence[np.ndarray], placements: dict[int, tuple[int, int]]) -> Image.Image: (empty)

16 if not placements: 16 ↛ 17line 16 didn't jump to line 17 because the condition on line 16 was never true(empty)

17 raise ValueError("No placements available") 

18 

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

20 xs = [coord[0] for coord in placements.values()] (empty)

21 ys = [coord[1] for coord in placements.values()] (empty)

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

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

24 

25 canvas_w = (maxx - minx + 1) * w (empty)

26 canvas_h = (maxy - miny + 1) * h (empty)

27 canvas = np.ones((canvas_h, canvas_w, 3), dtype=np.float32) (empty)

28 

29 for idx, (gx, gy) in placements.items(): (empty)

30 x = (gx - minx) * w (empty)

31 y = (maxy - gy) * h (empty)

32 canvas[y : y + h, x : x + w, :] = tiles[idx] (empty)

33 

34 return Image.fromarray(np.clip(canvas * 255.0, 0, 255).astype(np.uint8), mode="RGB") (empty)

35 

36 

37def tile_rgba_images(tiles: Sequence[np.ndarray]) -> List[Image.Image]: (empty)

38 out: List[Image.Image] = [] (empty)

39 for tile in tiles: (empty)

40 rgb = np.clip(tile * 255.0, 0, 255).astype(np.uint8) (empty)

41 out.append(Image.fromarray(rgb, mode="RGB").convert("RGBA")) (empty)

42 return out (empty)

43 

44 

45def component_image(tile_imgs: Sequence[Image.Image], adjs: AdjList, component: Sequence[int]) -> Image.Image: (empty)

46 root = component[0] (empty)

47 comp_set = set(component) (empty)

48 coords = {node: coord for node, coord in chargeds(adjs, root) if node in comp_set} (empty)

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

50 

51 xs = [c[0] for c in coords.values()] (empty)

52 ys = [c[1] for c in coords.values()] (empty)

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

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

55 

56 out_w = (maxx - minx + 1) * tile_w (empty)

57 out_h = (maxy - miny + 1) * tile_h (empty)

58 out = Image.new("RGBA", (out_w, out_h), (255, 255, 255, 0)) (empty)

59 

60 for node, (x, y) in coords.items(): (empty)

61 px = (x - minx) * tile_w (empty)

62 py = (maxy - y) * tile_h (empty)

63 out.alpha_composite(tile_imgs[node], (px, py)) (empty)

64 return out (empty)

65 

66 

67def boxes_overlap(a: Tuple[int, int, int, int], b: Tuple[int, int, int, int], margin: int) -> bool: (empty)

68 ax1, ay1, ax2, ay2 = a (empty)

69 bx1, by1, bx2, by2 = b (empty)

70 return not (ax2 + margin <= bx1 or bx2 + margin <= ax1 or ay2 + margin <= by1 or by2 + margin <= ay1) (empty)

71 

72 

73def pack_images_non_overlapping( (empty)

74 images: Sequence[Image.Image], 

75 canvas_size: Tuple[int, int], 

76 rng: random.Random, 

77 margin: int = 10, 

78 tries_per_image: int = 200, 

79) -> Image.Image: 

80 cw, ch = canvas_size (empty)

81 canvas = Image.new("RGBA", (cw, ch), (255, 255, 255, 255)) (empty)

82 placed_boxes: List[Tuple[int, int, int, int]] = [] (empty)

83 

84 for img in images: (empty)

85 placed = False (empty)

86 iw, ih = img.size (empty)

87 if iw > cw or ih > ch: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true(empty)

88 raise RuntimeError("Could not pack all components without overlap; try a larger animation frame size.") 

89 

90 for use_margin in range(margin, -1, -1): (empty)

91 max_x = cw - iw - use_margin (empty)

92 max_y = ch - ih - use_margin (empty)

93 if max_x < use_margin or max_y < use_margin: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true(empty)

94 continue 

95 

96 for _ in range(tries_per_image): (empty)

97 x = rng.randint(use_margin, max_x) (empty)

98 y = rng.randint(use_margin, max_y) (empty)

99 box = (x, y, x + iw, y + ih) (empty)

100 if any(boxes_overlap(box, old, use_margin) for old in placed_boxes): (empty)

101 continue (empty)

102 canvas.alpha_composite(img, (x, y)) (empty)

103 placed_boxes.append(box) (empty)

104 placed = True (empty)

105 break (empty)

106 if placed: (empty)

107 break (empty)

108 

109 stride = max(1, min(8, max(1, use_margin))) (empty)

110 for y in range(use_margin, max_y + 1, stride): (empty)

111 ok = False (empty)

112 for x in range(use_margin, max_x + 1, stride): (empty)

113 box = (x, y, x + iw, y + ih) (empty)

114 if any(boxes_overlap(box, old, use_margin) for old in placed_boxes): 114 ↛ 116line 114 didn't jump to line 116 because the condition on line 114 was always true(empty)

115 continue (empty)

116 canvas.alpha_composite(img, (x, y)) 

117 placed_boxes.append(box) 

118 placed = True 

119 ok = True 

120 break 

121 if ok: 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true(empty)

122 break 

123 if placed: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true(empty)

124 break 

125 

126 if not placed: (empty)

127 raise RuntimeError("Could not pack all components without overlap; try a larger animation frame size.") (empty)

128 

129 return canvas.convert("RGB") (empty)

130 

131 

132def minimum_canvas_side(images: Sequence[Image.Image], margin: int) -> int: (empty)

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

134 return 1 

135 

136 max_dim = max(max(img.size) for img in images) + 2 * margin (empty)

137 total_area = sum((img.size[0] + margin) * (img.size[1] + margin) for img in images) (empty)

138 return max(max_dim, int(math.ceil(math.sqrt(total_area)))) (empty)

139 

140 

141def pack_images_with_growing_canvas( (empty)

142 images: Sequence[Image.Image], 

143 rng: random.Random, 

144 frame_size: int, 

145 margin: int = 12, 

146 max_growth_steps: int = 12, 

147) -> Image.Image: 

148 work_size = max(frame_size, minimum_canvas_side(images, margin)) (empty)

149 for _ in range(max_growth_steps): 149 ↛ 154line 149 didn't jump to line 154 because the loop on line 149 didn't complete(empty)

150 try: (empty)

151 return pack_images_non_overlapping(images, (work_size, work_size), rng=rng, margin=margin) (empty)

152 except RuntimeError: (empty)

153 work_size = int(math.ceil(work_size * 1.25)) (empty)

154 raise RuntimeError("Failed to pack animation frame with all components.") 

155 

156 

157def pad_image_to_square(image: Image.Image, side: int) -> Image.Image: (empty)

158 if image.width > side or image.height > side: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true(empty)

159 raise ValueError("Image is larger than the requested square canvas.") 

160 

161 canvas = Image.new("RGB", (side, side), (255, 255, 255)) (empty)

162 x = (side - image.width) // 2 (empty)

163 y = (side - image.height) // 2 (empty)

164 canvas.paste(image, (x, y)) (empty)

165 return canvas (empty)

166 

167 

168def build_tree_animation_frames( (empty)

169 tiles: Sequence[np.ndarray], 

170 history: Sequence[AdjList], 

171 seed: int = 0, 

172 frame_size: int = 1024, 

173 max_angle: float = 35.0, 

174) -> List[Image.Image]: 

175 rng = random.Random(seed) (empty)

176 tile_imgs = tile_rgba_images(tiles) (empty)

177 frames: List[Image.Image] = [] (empty)

178 max_side = frame_size (empty)

179 

180 for adjs in history: (empty)

181 components = connected_components(adjs) (empty)

182 component_imgs: List[Image.Image] = [] (empty)

183 for comp in components: (empty)

184 img = component_image(tile_imgs, adjs, comp) (empty)

185 angle = rng.uniform(-max_angle, max_angle) (empty)

186 rotated = img.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC) (empty)

187 component_imgs.append(rotated) (empty)

188 

189 component_imgs.sort(key=lambda im: im.size[0] * im.size[1], reverse=True) (empty)

190 frame = pack_images_with_growing_canvas(component_imgs, rng=rng, frame_size=frame_size, margin=12) (empty)

191 max_side = max(max_side, frame.width, frame.height) (empty)

192 frames.append(frame) (empty)

193 

194 return [pad_image_to_square(frame, max_side) for frame in frames] (empty)

195 

196 

197def save_tree_build_animation( (empty)

198 tiles: Sequence[np.ndarray], 

199 history: Sequence[AdjList], 

200 output_path: Path, 

201 seed: int = 0, 

202 frame_size: int = 1024, 

203 max_angle: float = 35.0, 

204 duration_ms: int = 1000, 

205 frames_dir: Path | None = None, 

206) -> None: 

207 frames = build_tree_animation_frames( (empty)

208 tiles=tiles, 

209 history=history, 

210 seed=seed, 

211 frame_size=frame_size, 

212 max_angle=max_angle, 

213 ) 

214 if not frames: 214 ↛ 215line 214 didn't jump to line 215 because the condition on line 214 was never true(empty)

215 raise ValueError("No animation frames generated") 

216 final_placements = reconstruct_layout(history[-1]) (empty)

217 final_frame = render_reconstruction(tiles, final_placements) (empty)

218 common_side = max( (empty)

219 max(frame.width for frame in frames), 

220 max(frame.height for frame in frames), 

221 final_frame.width, 

222 final_frame.height, 

223 ) 

224 frames = [pad_image_to_square(frame, common_side) for frame in frames] (empty)

225 frames.append(pad_image_to_square(final_frame, common_side)) (empty)

226 if frames_dir is not None: 226 ↛ 232line 226 didn't jump to line 232 because the condition on line 226 was always true(empty)

227 frames_dir.mkdir(parents=True, exist_ok=True) (empty)

228 for existing in frames_dir.glob("frame_*.png"): (empty)

229 existing.unlink() (empty)

230 for i, frame in enumerate(frames): (empty)

231 frame.save(frames_dir / f"frame_{i:04d}.png") (empty)

232 frames[0].save( (empty)

233 output_path, 

234 save_all=True, 

235 append_images=frames[1:], 

236 duration=duration_ms, 

237 loop=0, 

238 optimize=False, 

239 )