Skip to content

Commit a3e5b77

Browse files
rudrankriyamclaude
andcommitted
Refactor k-means clustering and remove empty test file
- Add guard for empty pixels array in kMeansCluster to prevent crashes - Extract k-means logic into smaller helper functions for better cyclomatic complexity and readability - Remove empty LibraryPlaylistTests.swift file 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
1 parent dd67e37 commit a3e5b77

File tree

2 files changed

+54
-26
lines changed

2 files changed

+54
-26
lines changed

Sources/MusadoraKit/Extension/CommonImageProcessing.swift

Lines changed: 54 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -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.

Tests/MusadoraKitTests/LibraryPlaylistTests.swift

Whitespace-only changes.

0 commit comments

Comments
 (0)