Coverage for imcluster/features.py: 100.00%

1"""Feature extraction using pretrained Hugging Face vision models.""" 

2 

3from enum import Enum 

4from pathlib import Path 

5from typing import Any 

6 

7import numpy as np 

8import torch 

9from huggingface_hub import hf_hub_download, snapshot_download 

10from huggingface_hub.errors import HfHubHTTPError 

11from numpy.typing import NDArray 

12from PIL import Image, ImageOps, UnidentifiedImageError 

13from rich.console import Console 

14from rich.progress import track 

15from transformers import pipeline 

16 

17from .io import ImclusterIO 

18 

19console = Console() 

20 

21 

22class ModelArchitecture(str, Enum): 

23 """Supported DINOv3 architecture families.""" 

24 

25 VIT = "vit" 

26 CONVNEXT = "convnext" 

27 

28 

29class DinoVersion(str, Enum): 

30 """Supported DINO model generations.""" 

31 

32 AUTO = "auto" 

33 TWO = "2" 

34 THREE = "3" 

35 

36 

37class ModelSize(str, Enum): 

38 """Supported model-size tiers.""" 

39 

40 TINY = "tiny" 

41 SMALL = "small" 

42 BASE = "base" 

43 LARGE = "large" 

44 HUGE = "huge" 

45 MAX = "max" 

46 

47 

48class Device(str, Enum): 

49 """Supported inference devices.""" 

50 

51 AUTO = "auto" 

52 CPU = "cpu" 

53 CUDA = "cuda" 

54 MPS = "mps" 

55 

56 

57VIT_MODELS = { 

58 ModelSize.TINY: "facebook/dinov3-vits16-pretrain-lvd1689m", 

59 ModelSize.SMALL: "facebook/dinov3-vits16plus-pretrain-lvd1689m", 

60 ModelSize.BASE: "facebook/dinov3-vitb16-pretrain-lvd1689m", 

61 ModelSize.LARGE: "facebook/dinov3-vitl16-pretrain-lvd1689m", 

62 ModelSize.HUGE: "facebook/dinov3-vith16plus-pretrain-lvd1689m", 

63 ModelSize.MAX: "facebook/dinov3-vit7b16-pretrain-lvd1689m", 

64} 

65 

66CONVNEXT_MODELS = { 

67 ModelSize.TINY: "facebook/dinov3-convnext-tiny-pretrain-lvd1689m", 

68 ModelSize.SMALL: "facebook/dinov3-convnext-small-pretrain-lvd1689m", 

69 ModelSize.BASE: "facebook/dinov3-convnext-base-pretrain-lvd1689m", 

70 ModelSize.LARGE: "facebook/dinov3-convnext-large-pretrain-lvd1689m", 

71} 

72 

73DINOV2_MODELS = { 

74 ModelSize.SMALL: "facebook/dinov2-small", 

75 ModelSize.BASE: "facebook/dinov2-base", 

76 ModelSize.LARGE: "facebook/dinov2-large", 

77 ModelSize.MAX: "facebook/dinov2-giant", 

78} 

79 

80DINOV2_FALLBACK_MODELS = { 

81 ModelSize.TINY: DINOV2_MODELS[ModelSize.SMALL], 

82 ModelSize.SMALL: DINOV2_MODELS[ModelSize.SMALL], 

83 ModelSize.BASE: DINOV2_MODELS[ModelSize.BASE], 

84 ModelSize.LARGE: DINOV2_MODELS[ModelSize.LARGE], 

85 ModelSize.HUGE: DINOV2_MODELS[ModelSize.MAX], 

86 ModelSize.MAX: DINOV2_MODELS[ModelSize.MAX], 

87} 

88 

89DEFAULT_MODEL = DINOV2_MODELS[ModelSize.BASE] 

90DINOV3_ACCESS_DOCS = ( 

91 "https://rbturnbull.github.io/imcluster/models.html#getting-access-to-dinov3" 

92) 

93 

94 

95def dinov3_available(model_name: str) -> bool: 

96 """Return whether a DINOv3 model is cached or accessible from the Hub.""" 

97 try: 

98 snapshot = Path(snapshot_download(model_name, local_files_only=True)) 

99 if any(snapshot.glob("*.safetensors")) or any( 

100 snapshot.glob("pytorch_model*.bin") 

101 ): 

102 return True 

103 except (HfHubHTTPError, OSError): 

104 pass 

105 

106 try: 

107 hf_hub_download(model_name, "config.json") 

108 except (HfHubHTTPError, OSError): 

109 return False 

110 return True 

111 

112 

113def resolve_device(device: Device) -> str: 

114 """Resolve automatic inference device selection.""" 

115 if device is not Device.AUTO: 

116 return device.value 

117 if torch.cuda.is_available(): 

118 return Device.CUDA.value 

119 if torch.backends.mps.is_available(): 

120 return Device.MPS.value 

121 return Device.CPU.value 

122 

123 

124def resolve_model( 

125 model: str | None, 

126 dino_version: DinoVersion, 

127 architecture: ModelArchitecture, 

128 size: ModelSize, 

129) -> str: 

130 """Resolve a custom model or DINO architecture-size preset. 

131 

132 Args: 

133 model: Explicit Hugging Face model ID, which takes precedence when set. 

134 dino_version: DINO model generation used for preset selection. 

135 architecture: ViT or ConvNeXt architecture family. 

136 size: Requested model-size tier. 

137 

138 Returns: 

139 A Hugging Face model identifier. 

140 

141 Raises: 

142 ValueError: If the requested size is unavailable for the architecture. 

143 """ 

144 if model: 

145 return model 

146 

147 if dino_version is DinoVersion.TWO: 

148 models = DINOV2_MODELS 

149 else: 

150 models = ( 

151 VIT_MODELS if architecture is ModelArchitecture.VIT else CONVNEXT_MODELS 

152 ) 

153 try: 

154 selected_model = models[size] 

155 except KeyError as error: 

156 if dino_version is DinoVersion.AUTO: 

157 fallback = DINOV2_FALLBACK_MODELS[size] 

158 console.print( 

159 f"DINOv3 has no {architecture.value}/{size.value} preset; " 

160 f"using {fallback}" 

161 ) 

162 return fallback 

163 selection = ( 

164 "DINOv2" 

165 if dino_version is DinoVersion.TWO 

166 else f"DINOv3 architecture '{architecture.value}'" 

167 ) 

168 raise ValueError( 

169 f"Size '{size.value}' is not available for {selection}" 

170 ) from error 

171 

172 if dino_version is DinoVersion.AUTO: 

173 with console.status("[cyan]Checking DINOv3 model access...[/cyan]"): 

174 available = dinov3_available(selected_model) 

175 else: 

176 available = True 

177 if not available: 

178 fallback = DINOV2_FALLBACK_MODELS[size] 

179 console.print( 

180 "[bold yellow]Warning:[/bold yellow] DINOv3 is not available to " 

181 "the active Hugging Face account or local cache, so imcluster is " 

182 f"falling back to {fallback}. DINOv3 requires access approval and " 

183 "authentication. See the " 

184 f"[link={DINOV3_ACCESS_DOCS}]imcluster DINOv3 access instructions[/link]." 

185 ) 

186 return fallback 

187 return selected_model 

188 

189 

190def build_features( 

191 imcluster_io: ImclusterIO, 

192 model_name: str = DEFAULT_MODEL, 

193 device: Device = Device.AUTO, 

194 batch_size: int = 8, 

195 force: bool = False, 

196) -> NDArray[Any]: 

197 """Build or load normalized image feature vectors. 

198 

199 Args: 

200 imcluster_io: Image collection and its persisted result table. 

201 model_name: Hugging Face image-feature-extraction model identifier. 

202 device: Device used for model inference. 

203 batch_size: Number of images submitted for each inference call. 

204 force: Rebuild vectors even when a cached model column exists. 

205 

206 Returns: 

207 A two-dimensional array with one normalized feature vector per image. 

208 

209 Notes: 

210 Results are cached in a column named after ``model_name``. 

211 """ 

212 model_name = str(model_name) 

213 if batch_size < 1: 

214 raise ValueError("batch_size must be at least 1") 

215 

216 if not imcluster_io.has_column(model_name) or force: 

217 resolved_device = resolve_device(device) 

218 console.print( 

219 "[cyan]Loading feature model:[/cyan] " 

220 f"{model_name} on {resolved_device}; processing " 

221 f"{len(imcluster_io.images)} images in batches of {batch_size}." 

222 ) 

223 with console.status(f"[cyan]Loading model '{model_name}'...[/cyan]"): 

224 feature_extractor = pipeline( 

225 model=model_name, 

226 task="image-feature-extraction", 

227 device=resolved_device, 

228 ) 

229 

230 results: list[NDArray[Any]] = [] 

231 image_batches = [ 

232 imcluster_io.images[index : index + batch_size] 

233 for index in range(0, len(imcluster_io.images), batch_size) 

234 ] 

235 for paths in track(image_batches, description="Generating feature vectors:"): 

236 images = [] 

237 for path in paths: 

238 try: 

239 with Image.open(path) as image: 

240 images.append( 

241 ImageOps.exif_transpose(image).convert("RGB").copy() 

242 ) 

243 except (OSError, UnidentifiedImageError) as error: 

244 raise ValueError(f"Cannot read image '{path}': {error}") from error 

245 

246 outputs = feature_extractor(images, pool=True, batch_size=batch_size) 

247 for output in outputs: 

248 result = np.asarray(output) 

249 while result.ndim > 1 and result.shape[0] == 1: 

250 result = result[0] 

251 if result.ndim != 1: 

252 raise ValueError( 

253 f"Model '{model_name}' did not return pooled image embeddings" 

254 ) 

255 results.append(result) 

256 

257 feature_vectors = np.vstack(results) 

258 norms = np.linalg.norm( 

259 feature_vectors, 

260 axis=1, 

261 keepdims=True, 

262 ) 

263 if np.any(norms == 0): 

264 raise ValueError(f"Model '{model_name}' returned a zero-length embedding") 

265 feature_vectors /= norms 

266 

267 imcluster_io.save_column( 

268 model_name, [feature_vectors[x] for x in range(feature_vectors.shape[0])] 

269 ) 

270 console.print( 

271 "[green]Cached feature vectors:[/green] " 

272 f"wrote {len(feature_vectors)} embeddings to '{imcluster_io.output}'." 

273 ) 

274 else: 

275 console.print( 

276 "[green]Using cached feature vectors:[/green] " 

277 f"loaded model '{model_name}' embeddings from '{imcluster_io.output}'." 

278 ) 

279 

280 feature_vectors = np.array(imcluster_io.get_column(model_name).to_list()) 

281 

282 return feature_vectors