Fix startup crash: replace sklearn k-means with pure numpy implementation
sklearn was not installed in the container, causing ModuleNotFoundError on import of ai_tools.py and preventing the server from starting. Replaced with a self-contained numpy k-means++ implementation: - k-means++ seeding for better initial centers - 20-iteration Lloyd's algorithm - Same output: hex colors sorted by cluster frequency No new dependencies required. https://claude.ai/code/session_01B58MaJCU1R6KwBDJCp8AfN
This commit is contained in:
@@ -511,8 +511,6 @@ async def enhance(req: EnhanceRequest):
|
||||
|
||||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||||
|
||||
from sklearn.cluster import KMeans as _KMeans
|
||||
|
||||
class ExtractColorsRequest(BaseModel):
|
||||
image: str # base64
|
||||
count: int = 6
|
||||
@@ -520,28 +518,49 @@ class ExtractColorsRequest(BaseModel):
|
||||
|
||||
def _extract_colors(image_bytes: bytes, count: int) -> list[str]:
|
||||
"""
|
||||
Resize image to 150×150, k-means cluster pixels into `count` groups,
|
||||
sort by cluster size (largest first), return as hex strings.
|
||||
Resize image to 150×150, k-means cluster pixels into `count` groups
|
||||
using pure numpy (no sklearn dependency), return hex strings by frequency.
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
|
||||
count = max(1, min(count, 32))
|
||||
|
||||
pil = _Image.open(_io.BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||||
pixels = _np.array(pil, dtype=_np.float32).reshape(-1, 3) # (N, 3)
|
||||
pil = Image.open(BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||||
pixels = np.array(pil, dtype=np.float32).reshape(-1, 3) # (22500, 3)
|
||||
n = len(pixels)
|
||||
|
||||
km = _KMeans(n_clusters=count, n_init=10, random_state=42)
|
||||
labels = km.fit_predict(pixels)
|
||||
centers = km.cluster_centers_ # (count, 3)
|
||||
# Initialise centers with k-means++ seeding
|
||||
rng = np.random.default_rng(42)
|
||||
centers = [pixels[rng.integers(n)]]
|
||||
for _ in range(count - 1):
|
||||
dists = np.min([np.sum((pixels - c) ** 2, axis=1) for c in centers], axis=0)
|
||||
probs = dists / dists.sum()
|
||||
centers.append(pixels[rng.choice(n, p=probs)])
|
||||
centers = np.array(centers)
|
||||
|
||||
# Count pixels per cluster and sort by frequency descending
|
||||
counts = _np.bincount(labels, minlength=count)
|
||||
order = _np.argsort(-counts) # descending
|
||||
labels = np.zeros(n, dtype=np.int32)
|
||||
for _ in range(20): # max 20 iterations
|
||||
# Assign each pixel to nearest center
|
||||
dists = np.sum((pixels[:, None] - centers[None]) ** 2, axis=2) # (n, k)
|
||||
new_labels = np.argmin(dists, axis=1)
|
||||
if np.all(new_labels == labels):
|
||||
break
|
||||
labels = new_labels
|
||||
# Recompute centers
|
||||
for k in range(count):
|
||||
mask = labels == k
|
||||
if mask.any():
|
||||
centers[k] = pixels[mask].mean(axis=0)
|
||||
|
||||
hex_colors = []
|
||||
for idx in order:
|
||||
r, g, b = centers[idx].astype(int).clip(0, 255)
|
||||
hex_colors.append(f"#{r:02x}{g:02x}{b:02x}")
|
||||
counts = np.bincount(labels, minlength=count)
|
||||
order = np.argsort(-counts)
|
||||
|
||||
return hex_colors
|
||||
return [
|
||||
"#{:02x}{:02x}{:02x}".format(*centers[i].astype(int).clip(0, 255))
|
||||
for i in order
|
||||
]
|
||||
|
||||
|
||||
@router.post("/extract-colors")
|
||||
|
||||
Reference in New Issue
Block a user