Coverage for src / puzzletree / reconstruct / io.py: 92.11%

1from __future__ import annotations (empty)

2 

3from dataclasses import dataclass (empty)

4from pathlib import Path (empty)

5from typing import List (empty)

6 

7import numpy as np (empty)

8from PIL import Image (empty)

9 

10 

11@dataclass (empty)

12class TileWriteResult: (empty)

13 output_dir: Path (empty)

14 paths: List[Path] (empty)

15 rows: int (empty)

16 cols: int (empty)

17 tile_width: int (empty)

18 tile_height: int (empty)

19 

20 

21def to_float_array(img: Image.Image) -> np.ndarray: (empty)

22 arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0 (empty)

23 return arr (empty)

24 

25 

26def load_tiles_from_dir(input_dir: Path) -> List[np.ndarray]: (empty)

27 files = sorted([p for p in input_dir.iterdir() if p.suffix.lower() in {".bmp", ".png", ".jpg", ".jpeg"}]) (empty)

28 if not files: 28 ↛ 29line 28 didn't jump to line 29 because the condition on line 28 was never true(empty)

29 raise FileNotFoundError(f"No image tiles found in {input_dir}") 

30 

31 tiles = [to_float_array(Image.open(path)) for path in files] (empty)

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

33 for i, t in enumerate(tiles): (empty)

34 if t.shape[:2] != (h, w): 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true(empty)

35 raise ValueError(f"Tile {files[i]} has shape {t.shape[:2]}, expected {(h, w)}") 

36 return tiles (empty)

37 

38 

39def split_image_into_tiles(img: Image.Image, rows: int, cols: int) -> List[Image.Image]: (empty)

40 if rows <= 0 or cols <= 0: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true(empty)

41 raise ValueError("rows and cols must both be positive integers") 

42 

43 source = img.convert("RGB") (empty)

44 if source.height % rows != 0 or source.width % cols != 0: (empty)

45 raise ValueError( (empty)

46 f"Image size {(source.width, source.height)} is not evenly divisible by rows={rows}, cols={cols}.", 

47 ) 

48 

49 tile_h = source.height // rows (empty)

50 tile_w = source.width // cols (empty)

51 tiles: List[Image.Image] = [] (empty)

52 for row in range(rows): (empty)

53 for col in range(cols): (empty)

54 box = (col * tile_w, row * tile_h, (col + 1) * tile_w, (row + 1) * tile_h) (empty)

55 tiles.append(source.crop(box).copy()) (empty)

56 return tiles (empty)

57 

58 

59def save_tiles_from_image( (empty)

60 input_image: Path, 

61 output_dir: Path, 

62 rows: int, 

63 cols: int, 

64 prefix: str = "tile", 

65 overwrite: bool = False, 

66) -> TileWriteResult: 

67 existing_paths = sorted(output_dir.glob(f"{prefix}_*.png")) if output_dir.exists() else [] (empty)

68 if existing_paths and not overwrite: (empty)

69 raise FileExistsError( (empty)

70 f"Output directory {output_dir} already contains {prefix}_*.png files. Use --overwrite to replace them.", 

71 ) 

72 

73 output_dir.mkdir(parents=True, exist_ok=True) (empty)

74 for path in existing_paths: (empty)

75 path.unlink() (empty)

76 

77 with Image.open(input_image) as img: (empty)

78 tiles = split_image_into_tiles(img, rows=rows, cols=cols) (empty)

79 

80 written_paths: List[Path] = [] (empty)

81 for index, tile in enumerate(tiles): (empty)

82 path = output_dir / f"{prefix}_{index:03d}.png" (empty)

83 tile.save(path) (empty)

84 written_paths.append(path) (empty)

85 

86 return TileWriteResult( (empty)

87 output_dir=output_dir, 

88 paths=written_paths, 

89 rows=rows, 

90 cols=cols, 

91 tile_width=tiles[0].width, 

92 tile_height=tiles[0].height, 

93 )