@@ -143,44 +143,72 @@ struct CommonImageProcessing {
143143 ///
144144 /// - Returns: An array of Cluster objects representing the final clusters.
145145 private static func kMeansCluster( pixels: [ PixelData ] , k: Int , maxIterations: Int = 10 ) -> [ Cluster ] {
146+ guard !pixels. isEmpty else {
147+ return [ ]
148+ }
149+
150+ var clusters = initializeClusters ( pixels: pixels, k: k)
151+
152+ for _ in 0 ..< maxIterations {
153+ assignPixelsToClusters ( pixels: pixels, clusters: & clusters)
154+ updateClusterCenters ( clusters: & clusters)
155+ }
156+
157+ return clusters
158+ }
159+
160+ private static func initializeClusters( pixels: [ PixelData ] , k: Int ) -> [ Cluster ] {
146161 var clusters = [ Cluster] ( )
147162 for _ in 0 ..< k {
148163 if let randomPixel = pixels. randomElement ( ) {
149164 clusters. append ( Cluster ( center: randomPixel, points: [ ] ) )
150165 }
151166 }
167+ return clusters
168+ }
152169
153- for _ in 0 ..< maxIterations {
154- for clusterIndex in 0 ..< clusters. count {
155- clusters [ clusterIndex] . points. removeAll ( )
156- }
170+ private static func assignPixelsToClusters( pixels: [ PixelData ] , clusters: inout [ Cluster ] ) {
171+ clearClusterPoints ( clusters: & clusters)
157172
158- for pixel in pixels {
159- var minDistance = Double . greatestFiniteMagnitude
160- var closestClusterIndex = 0
161- for (index, cluster) in clusters. enumerated ( ) {
162- let distance = euclideanDistance ( pixel1: pixel, pixel2: cluster. center)
163- if distance < minDistance {
164- minDistance = distance
165- closestClusterIndex = index
166- }
167- }
168- clusters [ closestClusterIndex] . points. append ( pixel)
169- }
173+ for pixel in pixels {
174+ let closestIndex = findClosestClusterIndex ( pixel: pixel, clusters: clusters)
175+ clusters [ closestIndex] . points. append ( pixel)
176+ }
177+ }
170178
171- for clusterIndex in 0 ..< clusters. count {
172- let cluster = clusters [ clusterIndex]
173- guard !cluster. points. isEmpty else { continue }
174- let sum = cluster. points. reduce ( PixelData ( red: 0 , green: 0 , blue: 0 ) ) { result, pixel -> PixelData in
175- return PixelData ( red: result. red + pixel. red, green: result. green + pixel. green, blue: result. blue + pixel. blue)
176- }
177- let count = Double ( cluster. points. count)
178- guard count > 0 else { continue }
179- clusters [ clusterIndex] . center = PixelData ( red: sum. red / count, green: sum. green / count, blue: sum. blue / count)
179+ private static func clearClusterPoints( clusters: inout [ Cluster ] ) {
180+ for index in 0 ..< clusters. count {
181+ clusters [ index] . points. removeAll ( )
182+ }
183+ }
184+
185+ private static func findClosestClusterIndex( pixel: PixelData , clusters: [ Cluster ] ) -> Int {
186+ var minDistance = Double . greatestFiniteMagnitude
187+ var closestIndex = 0
188+
189+ for (index, cluster) in clusters. enumerated ( ) {
190+ let distance = euclideanDistance ( pixel1: pixel, pixel2: cluster. center)
191+ if distance < minDistance {
192+ minDistance = distance
193+ closestIndex = index
180194 }
181195 }
182196
183- return clusters
197+ return closestIndex
198+ }
199+
200+ private static func updateClusterCenters( clusters: inout [ Cluster ] ) {
201+ for index in 0 ..< clusters. count {
202+ let cluster = clusters [ index]
203+ guard !cluster. points. isEmpty else { continue }
204+
205+ let sum = cluster. points. reduce ( PixelData ( red: 0 , green: 0 , blue: 0 ) ) { result, pixel in
206+ PixelData ( red: result. red + pixel. red, green: result. green + pixel. green, blue: result. blue + pixel. blue)
207+ }
208+
209+ let count = Double ( cluster. points. count)
210+ clusters [ index] . center = PixelData ( red: sum. red / count, green: sum. green / count, blue: sum. blue / count)
211+ }
184212 }
185213
186214 /// Calculates the Euclidean distance between two pixels in color space.
0 commit comments