Coverage for imcluster/main.py: 100.00%
95 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-12 01:58 +0000
1"""Command-line entry point for the imcluster pipeline."""
3import tempfile
4import webbrowser
5from pathlib import Path
6from typing import Annotated
8import typer
9from rich.console import Console
11from .cluster import ClusteringAlgorithm, cluster
12from .evaluate import evaluate_clustering, print_evaluation, write_evaluation
13from .features import (
14 Device,
15 DinoVersion,
16 ModelArchitecture,
17 ModelSize,
18 build_features,
19 resolve_model,
20)
21from .html import write_html
22from .io import ImclusterIO
23from .llm import DEFAULT_LLM, name_clusters
24from .reduction import ReductionMethod, reduce_dimensions
25from .thumbnails import generate_thumbnails
27console = Console()
29app = typer.Typer()
32def open_gallery(path: Path) -> None:
33 """Open a generated HTML gallery in the default web browser."""
34 webbrowser.open(path.resolve().as_uri())
37@app.command()
38def main(
39 inputs: Annotated[
40 list[Path] | None,
41 typer.Argument(
42 help=(
43 "Image files, directories, or text manifests. May be omitted when "
44 "loading an existing --cache file."
45 )
46 ),
47 ] = None,
48 cache: Annotated[
49 Path | None,
50 typer.Option(help="Preserve processing results in this Parquet file."),
51 ] = None,
52 gallery: Annotated[
53 Path | None,
54 typer.Option(help="Preserve the HTML gallery at this path."),
55 ] = None,
56 dino_version: Annotated[
57 DinoVersion,
58 typer.Option(help="DINO model generation used for preset selection."),
59 ] = DinoVersion.AUTO,
60 arch: Annotated[
61 ModelArchitecture,
62 typer.Option(help="DINOv3 architecture family; ignored for DINOv2."),
63 ] = ModelArchitecture.VIT,
64 size: Annotated[
65 ModelSize,
66 typer.Option(help="DINO model size used when --model is not supplied."),
67 ] = ModelSize.BASE,
68 model: Annotated[
69 str | None,
70 typer.Option(
71 help="Hugging Face model ID; overrides --dino-version, --arch, and --size."
72 ),
73 ] = None,
74 device: Annotated[
75 Device,
76 typer.Option(help="Device for model inference."),
77 ] = Device.AUTO,
78 batch_size: Annotated[
79 int,
80 typer.Option(
81 min=1,
82 help="Number of images processed per inference batch.",
83 ),
84 ] = 8,
85 max_images: Annotated[
86 int | None,
87 typer.Option(min=1, help="Maximum number of input images to process."),
88 ] = None,
89 recursive: Annotated[
90 bool,
91 typer.Option(help="Search directory inputs recursively."),
92 ] = False,
93 clustering: Annotated[
94 ClusteringAlgorithm,
95 typer.Option(help="Clustering algorithm to use."),
96 ] = ClusteringAlgorithm.KMEANS,
97 reduce: Annotated[
98 ReductionMethod,
99 typer.Option(help="Dimensionality reduction applied before clustering."),
100 ] = ReductionMethod.UMAP,
101 reduction_dims: Annotated[
102 int,
103 typer.Option(min=1, help="Target number of dimensions after reduction."),
104 ] = 50,
105 n_clusters: Annotated[
106 int,
107 typer.Option(min=2, help="Number of clusters for fixed-count methods."),
108 ] = 20,
109 dbscan_eps: Annotated[
110 float,
111 typer.Option(help="Maximum cosine distance between DBSCAN neighbours."),
112 ] = 0.5,
113 min_samples: Annotated[
114 int,
115 typer.Option(min=1, help="Minimum sample count for density clustering."),
116 ] = 2,
117 name: Annotated[
118 bool,
119 typer.Option(help="Generate descriptive cluster names with a multimodal LLM."),
120 ] = False,
121 llm: Annotated[
122 str,
123 typer.Option(help="llmloader model identifier used to name clusters."),
124 ] = DEFAULT_LLM,
125 llm_temperature: Annotated[
126 float,
127 typer.Option(min=0.0, help="Sampling temperature used for cluster names."),
128 ] = 0.2,
129 llm_api_key: Annotated[
130 str | None,
131 typer.Option(help="Optional provider API key used by the naming LLM."),
132 ] = None,
133 in_group_size: Annotated[
134 int,
135 typer.Option(
136 min=1,
137 help="Maximum in-cluster thumbnail examples sent to the naming LLM.",
138 ),
139 ] = 10,
140 out_group_size: Annotated[
141 int,
142 typer.Option(
143 min=0,
144 help="Maximum contrasting outside examples sent to the naming LLM.",
145 ),
146 ] = 0,
147 evaluate: Annotated[
148 Path | None,
149 typer.Option(
150 "--evaluate",
151 "--expected",
152 exists=True,
153 dir_okay=False,
154 readable=True,
155 help="CSV containing filename,class labels for clustering evaluation.",
156 ),
157 ] = None,
158 metric: Annotated[
159 Path | None,
160 typer.Option(help="Write NMI, ARI, and ACC evaluation scores to this CSV."),
161 ] = None,
162 thumbnail_width: Annotated[
163 int,
164 typer.Option(min=1, help="Maximum thumbnail width in pixels."),
165 ] = 256,
166 thumbnail_height: Annotated[
167 int,
168 typer.Option(min=1, help="Maximum thumbnail height in pixels."),
169 ] = 256,
170 force: Annotated[
171 bool,
172 typer.Option(help="Recompute all cached processing stages."),
173 ] = False,
174 force_features: Annotated[
175 bool,
176 typer.Option(help="Recompute feature vectors and downstream stages."),
177 ] = False,
178 force_cluster: Annotated[
179 bool,
180 typer.Option(help="Recompute cluster labels."),
181 ] = False,
182 force_thumbnails: Annotated[
183 bool,
184 typer.Option(help="Regenerate cached thumbnails."),
185 ] = False,
186 no_open: Annotated[
187 bool,
188 typer.Option("--no-open", help="Do not open the generated gallery."),
189 ] = False,
190) -> None:
191 """Cluster images and open an HTML gallery."""
192 if metric is not None and evaluate is None:
193 raise typer.BadParameter(
194 "--metric requires --evaluate expected_classes.csv",
195 param_hint="--metric",
196 )
198 temporary_directory: Path | None = None
199 if cache is None:
200 temporary_directory = Path(tempfile.mkdtemp(prefix="imcluster-"))
201 output_df = temporary_directory / "results.parquet"
202 else:
203 output_df = cache
204 if gallery is None:
205 if temporary_directory is None:
206 temporary_directory = Path(tempfile.mkdtemp(prefix="imcluster-"))
207 output_html = temporary_directory / "gallery.html"
208 else:
209 output_html = gallery
211 try:
212 imcluster_io = ImclusterIO(
213 inputs or [],
214 output_df,
215 max_images=max_images,
216 recursive=recursive,
217 reset_cache=force,
218 )
219 except ValueError as error:
220 raise typer.BadParameter(str(error), param_hint="--cache") from error
221 cached_model_names = (
222 {
223 value
224 for value in imcluster_io.df["model"].tolist()
225 if isinstance(value, str) and value
226 }
227 if not inputs and model is None and imcluster_io.has_column("model")
228 else set()
229 )
230 cached_model_name = next(iter(cached_model_names), None)
231 if (
232 len(cached_model_names) == 1
233 and cached_model_name is not None
234 and imcluster_io.has_column(cached_model_name)
235 and not force_features
236 ):
237 model_name = cached_model_name
238 console.print(
239 f"[green]Using cached model:[/green] restored '{model_name}' from "
240 f"'{imcluster_io.output}'."
241 )
242 else:
243 try:
244 model_name = resolve_model(model, dino_version, arch, size)
245 except ValueError as error:
246 raise typer.BadParameter(str(error), param_hint="--size") from error
248 if not imcluster_io.images:
249 raise typer.BadParameter(
250 "No valid input images were found. Provide image inputs or an existing "
251 "--cache file.",
252 param_hint="inputs",
253 )
254 if len(imcluster_io.images) < 2:
255 raise typer.BadParameter(
256 "At least two images are required", param_hint="inputs"
257 )
259 console.print(
260 f"[bold]Processing {len(imcluster_io.images)} images[/bold] with model "
261 f"'{model_name}' and {clustering.value} clustering."
262 )
264 feature_vectors = build_features(
265 imcluster_io,
266 model_name=model_name,
267 device=device,
268 batch_size=batch_size,
269 force=force or force_features,
270 )
271 clustering_vectors = reduce_dimensions(
272 imcluster_io,
273 feature_vectors,
274 method=reduce,
275 dimensions=reduction_dims,
276 force=force or force_features,
277 )
278 previous_reductions = (
279 {
280 value
281 for value in imcluster_io.df["reduction"].tolist()
282 if isinstance(value, str) and value
283 }
284 if imcluster_io.has_column("reduction")
285 else {ReductionMethod.NONE.value}
286 )
287 reduction_changed = previous_reductions != {reduce.value}
288 previous_reduction_dims = (
289 set(imcluster_io.df["reduction_dims"].dropna().astype(int).tolist())
290 if imcluster_io.has_column("reduction_dims")
291 else set()
292 )
293 reduction_dims_changed = (
294 reduce is not ReductionMethod.NONE
295 and previous_reduction_dims != {reduction_dims}
296 )
297 cluster_column = f"{clustering.value}_cluster"
298 cluster_force = (
299 force
300 or force_features
301 or force_cluster
302 or reduction_changed
303 or reduction_dims_changed
304 )
305 cluster_was_recomputed = cluster_force or not imcluster_io.has_column(
306 cluster_column
307 )
308 cluster(
309 imcluster_io,
310 clustering_vectors,
311 algorithm=clustering,
312 n_clusters=n_clusters,
313 dbscan_eps=dbscan_eps,
314 min_samples=min_samples,
315 force=cluster_force,
316 )
317 cluster_name_column = f"{cluster_column}_name"
318 if cluster_was_recomputed and imcluster_io.has_column(cluster_name_column):
319 imcluster_io.df.drop(columns=[cluster_name_column], inplace=True)
320 imcluster_io.df["model"] = model_name
321 imcluster_io.df["algorithm"] = clustering.value
322 imcluster_io.df["reduction"] = reduce.value
323 imcluster_io.df["reduction_dims"] = reduction_dims
324 imcluster_io.save()
325 if evaluate is not None:
326 try:
327 metrics = evaluate_clustering(
328 imcluster_io,
329 evaluate,
330 cluster_column=f"{clustering.value}_cluster",
331 )
332 except ValueError as error:
333 raise typer.BadParameter(str(error), param_hint="--evaluate") from error
334 print_evaluation(metrics)
335 if metric is not None:
336 write_evaluation(metrics, metric)
337 console.print(
338 f"[green]Wrote evaluation metrics:[/green] {metric.resolve()}"
339 )
340 generate_thumbnails(
341 imcluster_io,
342 thumbnail_height=thumbnail_height,
343 thumbnail_width=thumbnail_width,
344 force=force,
345 force_thumbnails=force_thumbnails,
346 )
347 if name:
348 try:
349 name_clusters(
350 imcluster_io,
351 feature_vectors,
352 cluster_column=cluster_column,
353 llm=llm,
354 temperature=llm_temperature,
355 api_key=llm_api_key,
356 in_group_size=in_group_size,
357 out_group_size=out_group_size,
358 force=force or cluster_was_recomputed,
359 )
360 except ValueError as error:
361 raise typer.BadParameter(str(error), param_hint="--name") from error
362 with console.status("[cyan]Rendering HTML gallery...[/cyan]"):
363 write_html(
364 imcluster_io,
365 output_html=output_html,
366 cluster_column=cluster_column,
367 metadata={
368 "Model": model_name,
369 "Clustering": clustering.value,
370 "Reduction": reduce.value,
371 "Reduction dimensions": str(reduction_dims),
372 "Images": str(len(imcluster_io.images)),
373 },
374 feature_vectors=feature_vectors,
375 )
376 console.print(f"[green]Wrote processing cache:[/green] {output_df.resolve()}")
377 console.print(f"[green]Wrote HTML gallery:[/green] {output_html.resolve()}")
378 if cache is None:
379 console.print(
380 "[yellow]Cache is temporary:[/yellow] no persistent cache file was "
381 "requested. Use [bold]--cache PATH[/bold] to preserve it."
382 )
383 if gallery is None:
384 console.print(
385 "[yellow]Gallery is temporary:[/yellow] no persistent gallery file "
386 "was requested. Use [bold]--gallery PATH[/bold] to preserve it."
387 )
388 if not no_open:
389 console.print(f"[cyan]Opening gallery:[/cyan] {output_html.resolve()}")
390 open_gallery(output_html)
391 else:
392 console.print(
393 "[dim]Gallery was not opened because --no-open was specified.[/dim]"
394 )