Coverage for imcluster/io.py: 100.00%

1"""Image input discovery and Parquet-backed result persistence.""" 

2 

3from collections.abc import Iterable 

4from pathlib import Path 

5from typing import Any 

6 

7import pandas as pd 

8from rich.console import Console 

9 

10console = Console() 

11 

12SUPPORTED_IMAGE_SUFFIXES = { 

13 ".png", 

14 ".jpg", 

15 ".jpeg", 

16 ".tiff", 

17 ".tif", 

18 ".bmp", 

19 ".gif", 

20} 

21 

22 

23def valid_image(path: str | Path) -> bool: 

24 """Return whether a path is an existing image with a supported suffix.""" 

25 path = Path(path) 

26 return path.is_file() and path.suffix.lower() in SUPPORTED_IMAGE_SUFFIXES 

27 

28 

29class ImclusterIO: 

30 """Manage image inputs and cached tabular results. 

31 

32 Args: 

33 inputs: Image files, directories, or text manifests containing one path 

34 per line. 

35 output: Parquet file used to persist intermediate and final results. 

36 max_images: Optional maximum number of images to retain. Zero or 

37 ``None`` means unlimited. 

38 recursive: Search within nested subdirectories of directory inputs. 

39 reset_cache: Discard an existing cache instead of loading it. 

40 """ 

41 

42 def __init__( 

43 self, 

44 inputs: Iterable[str | Path], 

45 output: str | Path, 

46 max_images: int | None = None, 

47 recursive: bool = False, 

48 reset_cache: bool = False, 

49 ) -> None: 

50 """Initialize an image collection and load any cached results.""" 

51 self.output: Path = Path(output) 

52 self.output.parent.mkdir(parents=True, exist_ok=True) 

53 self.images: list[Path] 

54 self.filenames: list[str] 

55 self.paths: list[str] 

56 self.df: pd.DataFrame 

57 supplied_inputs = list(inputs) 

58 

59 if not supplied_inputs and self.output.exists() and not reset_cache: 

60 df = pd.read_parquet(self.output) 

61 if "path" not in df: 

62 raise ValueError( 

63 f"Cached results in '{self.output}' do not contain image paths." 

64 ) 

65 self.paths = df["path"].tolist() 

66 self.images = [Path(path) for path in self.paths] 

67 self.filenames = ( 

68 df["filenames"].tolist() 

69 if "filenames" in df 

70 else [image.name for image in self.images] 

71 ) 

72 self.df = df 

73 return 

74 

75 discovered: list[Path] = [] 

76 for path in supplied_inputs: 

77 path = Path(path).expanduser() 

78 

79 if path.is_dir(): 

80 candidates = path.rglob("*") if recursive else path.iterdir() 

81 discovered.extend(sorted(x for x in candidates if valid_image(x))) 

82 elif path.suffix.lower() == ".txt": 

83 with open(path, encoding="utf-8") as f: 

84 for line in f: 

85 value = line.strip() 

86 if not value: 

87 continue 

88 candidate = Path(value).expanduser() 

89 if not candidate.is_absolute(): 

90 candidate = path.parent / candidate 

91 if valid_image(candidate): 

92 discovered.append(candidate) 

93 elif valid_image(path): 

94 discovered.append(path) 

95 else: 

96 console.print( 

97 "[yellow]Skipping input:[/yellow] " 

98 f"'{path}' is not an existing image with a supported extension." 

99 ) 

100 

101 # Resolve paths and remove duplicates without changing input order. 

102 self.images = list(dict.fromkeys(image.resolve() for image in discovered)) 

103 

104 if max_images and len(self.images) > max_images: 

105 self.images = self.images[:max_images] 

106 

107 self.filenames = [image.name for image in self.images] 

108 self.paths = [str(image) for image in self.images] 

109 

110 if self.output.exists() and not reset_cache: 

111 df = pd.read_parquet(self.output) 

112 if "path" not in df or df["path"].tolist() != self.paths: 

113 raise ValueError( 

114 f"Cached results in '{self.output}' do not match the current " 

115 "image inputs. Use --force to replace the cache." 

116 ) 

117 else: 

118 df = pd.DataFrame({"path": self.paths, "filenames": self.filenames}) 

119 

120 self.df = df 

121 

122 def has_column(self, column_name: str) -> bool: 

123 """Return whether the cached table contains a named column.""" 

124 return column_name in self.df.columns 

125 

126 def get_all_columns(self) -> list[str]: 

127 """Return all cached table column names.""" 

128 return self.df.columns.tolist() 

129 

130 def save(self) -> None: 

131 """Persist the current result table as Parquet.""" 

132 self.df.to_parquet(self.output, engine="pyarrow") 

133 

134 def save_column( 

135 self, 

136 column_name: str, 

137 data: Any, 

138 autosave: bool = True, 

139 ) -> None: 

140 """Add or replace a result column and optionally persist the table.""" 

141 self.df[column_name] = data 

142 if autosave: 

143 self.save() 

144 

145 def get_column(self, column_name: str) -> pd.Series: 

146 """Return a cached result column.""" 

147 return self.df[column_name]