Coverage for imcluster/cluster.py: 100.00%

1"""Clustering operations for image feature vectors.""" 

2 

3from enum import Enum 

4 

5import numpy as np 

6from numpy.typing import ArrayLike 

7from rich.console import Console 

8from sklearn.cluster import ( 

9 DBSCAN, 

10 HDBSCAN, 

11 AgglomerativeClustering, 

12 KMeans, 

13 SpectralClustering, 

14) 

15 

16from .io import ImclusterIO 

17 

18console = Console() 

19 

20 

21class ClusteringAlgorithm(str, Enum): 

22 """Supported clustering algorithms.""" 

23 

24 SPECTRAL = "spectral" 

25 DBSCAN = "dbscan" 

26 HDBSCAN = "hdbscan" 

27 KMEANS = "kmeans" 

28 AGGLOMERATIVE = "agglomerative" 

29 HIERARCHICAL = "hierarchical" 

30 

31 

32def cluster( 

33 imcluster_io: ImclusterIO, 

34 feature_vectors: ArrayLike, 

35 algorithm: ClusteringAlgorithm | str = ClusteringAlgorithm.SPECTRAL, 

36 n_clusters: int = 20, 

37 dbscan_eps: float = 0.5, 

38 min_samples: int = 2, 

39 force: bool = False, 

40) -> None: 

41 """Assign images to clusters and cache their labels. 

42 

43 Args: 

44 imcluster_io: Image collection and its persisted result table. 

45 feature_vectors: Feature matrix with one row per image. 

46 algorithm: Clustering algorithm to use. 

47 n_clusters: Number of clusters requested by fixed-count algorithms. 

48 dbscan_eps: Maximum cosine distance between DBSCAN neighbours. 

49 min_samples: Minimum neighbourhood or cluster size for density methods. 

50 force: Recompute labels even when the appropriate column already exists. 

51 

52 Raises: 

53 ValueError: If the algorithm or its parameters are invalid. 

54 """ 

55 try: 

56 algorithm = ( 

57 algorithm 

58 if isinstance(algorithm, ClusteringAlgorithm) 

59 else ClusteringAlgorithm(algorithm.lower()) 

60 ) 

61 except (AttributeError, ValueError) as error: 

62 raise ValueError(f"Unsupported clustering algorithm: {algorithm}") from error 

63 

64 vectors = np.asarray(feature_vectors) 

65 if vectors.ndim != 2 or len(vectors) != len(imcluster_io.images): 

66 raise ValueError("feature_vectors must contain one row per image") 

67 

68 cluster_column = f"{algorithm.value}_cluster" 

69 if imcluster_io.has_column(cluster_column) and not force: 

70 console.print( 

71 f"[green]Using cached {algorithm.value} clusters:[/green] " 

72 f"loaded {len(imcluster_io.images)} assignments from " 

73 f"'{imcluster_io.output}'." 

74 ) 

75 return 

76 

77 if min_samples < 1: 

78 raise ValueError("min_samples must be at least 1") 

79 

80 fixed_count = { 

81 ClusteringAlgorithm.SPECTRAL, 

82 ClusteringAlgorithm.KMEANS, 

83 ClusteringAlgorithm.AGGLOMERATIVE, 

84 ClusteringAlgorithm.HIERARCHICAL, 

85 } 

86 if algorithm in fixed_count: 

87 if n_clusters < 2: 

88 raise ValueError("n_clusters must be at least 2") 

89 if n_clusters > len(vectors): 

90 raise ValueError("n_clusters cannot exceed the number of images") 

91 

92 if algorithm is ClusteringAlgorithm.SPECTRAL: 

93 clustering = SpectralClustering(n_clusters=n_clusters, random_state=0) 

94 elif algorithm is ClusteringAlgorithm.KMEANS: 

95 clustering = KMeans(n_clusters=n_clusters, random_state=0, n_init="auto") 

96 elif algorithm is ClusteringAlgorithm.AGGLOMERATIVE: 

97 clustering = AgglomerativeClustering(n_clusters=n_clusters) 

98 elif algorithm is ClusteringAlgorithm.HIERARCHICAL: 

99 clustering = AgglomerativeClustering( 

100 n_clusters=n_clusters, 

101 metric="cosine", 

102 linkage="average", 

103 ) 

104 elif algorithm is ClusteringAlgorithm.HDBSCAN: 

105 if min_samples < 2: 

106 raise ValueError("min_samples must be at least 2 for HDBSCAN") 

107 clustering = HDBSCAN( 

108 min_cluster_size=min_samples, 

109 min_samples=min_samples, 

110 metric="cosine", 

111 algorithm="brute", 

112 copy=True, 

113 ) 

114 else: 

115 if dbscan_eps <= 0: 

116 raise ValueError("dbscan_eps must be greater than 0") 

117 clustering = DBSCAN( 

118 eps=dbscan_eps, 

119 min_samples=min_samples, 

120 metric="cosine", 

121 ) 

122 

123 console.print( 

124 f"[cyan]Running {algorithm.value} clustering:[/cyan] " 

125 f"assigning {len(imcluster_io.images)} images." 

126 ) 

127 with console.status( 

128 f"[cyan]Computing {algorithm.value} cluster assignments...[/cyan]" 

129 ): 

130 labels = clustering.fit_predict(vectors) 

131 imcluster_io.save_column(cluster_column, labels) 

132 console.print( 

133 f"[green]Cached {algorithm.value} clusters:[/green] " 

134 f"wrote {len(labels)} assignments to '{imcluster_io.output}'." 

135 )