Cut cluster scout from optimizer

src/repixelizer/analysis.py

@@ -55,91 +55,24 @@ def _compute_edge_map_torch(torch, rgba_t):
     return edge.to(dtype=torch.float32)
 
 
-def _kmeans(data: np.ndarray, k: int, seed: int, iterations: int = 12) -> tuple[np.ndarray, np.ndarray]:
-    rng = np.random.default_rng(seed)
-    if data.shape[0] <= k:
-        centers = data.copy()
-        labels = np.arange(data.shape[0], dtype=np.int32)
-        return centers, labels
-    picks = rng.choice(data.shape[0], size=k, replace=False)
-    centers = data[picks].copy()
-    for _ in range(iterations):
-        distances = np.linalg.norm(data[:, None, :] - centers[None, :, :], axis=-1)
-        labels = np.argmin(distances, axis=1)
-        new_centers = []
-        for idx in range(k):
-            members = data[labels == idx]
-            if members.size == 0:
-                new_centers.append(centers[idx])
-            else:
-                new_centers.append(np.mean(members, axis=0))
-        centers = np.asarray(new_centers, dtype=np.float32)
-    distances = np.linalg.norm(data[:, None, :] - centers[None, :, :], axis=-1)
-    labels = np.argmin(distances, axis=1)
-    return centers.astype(np.float32), labels.astype(np.int32)
-
-
-def _kmeans_torch(torch, data_t, k: int, seed: int, iterations: int = 12):
-    if data_t.shape[0] <= k:
-        centers_t = data_t.clone()
-        labels_t = torch.arange(data_t.shape[0], device=data_t.device, dtype=torch.long)
-        return centers_t, labels_t
-    generator = torch.Generator(device=data_t.device)
-    generator.manual_seed(seed)
-    picks = torch.randperm(data_t.shape[0], device=data_t.device, generator=generator)[:k]
-    centers_t = data_t[picks].clone()
-    for _ in range(iterations):
-        distances = torch.linalg.vector_norm(data_t[:, None, :] - centers_t[None, :, :], dim=-1)
-        labels_t = torch.argmin(distances, dim=1)
-        new_centers = []
-        for idx in range(k):
-            members = data_t[labels_t == idx]
-            if members.numel() == 0:
-                new_centers.append(centers_t[idx])
-            else:
-                new_centers.append(members.mean(dim=0))
-        centers_t = torch.stack(new_centers, dim=0).to(dtype=torch.float32)
-    distances = torch.linalg.vector_norm(data_t[:, None, :] - centers_t[None, :, :], dim=-1)
-    labels_t = torch.argmin(distances, dim=1)
-    return centers_t, labels_t
-
-
 def analyze_continuous_source(
     rgba: np.ndarray,
     seed: int,
-    cluster_count: int = 6,
     device: str | None = None,
 ) -> ContinuousSourceAnalysis:
+    del seed
     if device is not None:
         torch = _require_torch()
         resolved_device = _resolve_device(torch, device)
         rgba_t = torch.from_numpy(rgba).to(device=resolved_device, dtype=torch.float32)
         edge_map_t = _compute_edge_map_torch(torch, rgba_t)
-        alpha_t = rgba_t[..., 3]
-        cluster_map_t = torch.full(alpha_t.shape, -1, device=resolved_device, dtype=torch.long)
-        opaque_t = alpha_t > 0.05
-        if bool(torch.any(opaque_t).item()):
-            samples_t = rgba_t[opaque_t]
-            k = min(cluster_count, max(2, int(np.sqrt(int(samples_t.shape[0]) // 64 + 1))))
-            _centers_t, labels_t = _kmeans_torch(torch, samples_t, k=k, seed=seed)
-            cluster_map_t[opaque_t] = labels_t
         return ContinuousSourceAnalysis(
             edge_map=edge_map_t.detach().cpu().numpy().astype(np.float32),
-            cluster_map=cluster_map_t.detach().cpu().numpy().astype(np.int32),
         )
 
     edge_map = _compute_edge_map(rgba)
-    alpha = rgba[..., 3]
-    cluster_map = np.full(alpha.shape, -1, dtype=np.int32)
-    opaque = alpha > 0.05
-    if np.any(opaque):
-        samples = rgba[opaque]
-        k = min(cluster_count, max(2, int(np.sqrt(samples.shape[0] // 64 + 1))))
-        _centers, labels = _kmeans(samples, k=k, seed=seed)
-        cluster_map[opaque] = labels
     return ContinuousSourceAnalysis(
         edge_map=edge_map,
-        cluster_map=cluster_map,
     )
 
 

src/repixelizer/continuous.py

@@ -45,15 +45,6 @@ def _make_regular_uv(height: int, width: int, target_height: int, target_width:
     return uv.astype(np.float32)
 
 
-def _cluster_boundary_map(cluster_map: np.ndarray) -> np.ndarray:
-    if cluster_map.size == 0:
-        return np.zeros_like(cluster_map, dtype=np.float32)
-    boundary = np.zeros(cluster_map.shape, dtype=np.float32)
-    boundary[:, 1:] = np.maximum(boundary[:, 1:], (cluster_map[:, 1:] != cluster_map[:, :-1]).astype(np.float32))
-    boundary[1:, :] = np.maximum(boundary[1:, :], (cluster_map[1:, :] != cluster_map[:-1, :]).astype(np.float32))
-    return boundary
-
-
 def _edge_gradient_maps(edge_map: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     grad_x = np.zeros_like(edge_map, dtype=np.float32)
     grad_y = np.zeros_like(edge_map, dtype=np.float32)
@@ -109,16 +100,6 @@ def _representative_colors(patches, solver_params: SolverHyperParams):
     return representative, coherence
 
 
-def _exemplar_colors(patches):
-    patch_mean = patches.mean(dim=3, keepdim=True)
-    distances = (patches - patch_mean).abs().mean(dim=-1)
-    selected = distances.argmin(dim=3)
-    return patches.gather(
-        dim=3,
-        index=selected[..., None, None].expand(*selected.shape, 1, patches.shape[-1]),
-    ).squeeze(3)
-
-
 def _source_boundary_deltas(F, source_t, uv, axis: str, probe_scale: float = 0.22):
     if axis == "x":
         if uv.shape[2] < 2:
@@ -1355,7 +1336,7 @@ def _discrete_refine_output(
     return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history
 
 
-def optimize_uv_field(
+def optimize_lattice_pixels(
     rgba: np.ndarray,
     inference: InferenceResult,
     analysis: ContinuousSourceAnalysis,
@@ -1377,8 +1358,7 @@ def optimize_uv_field(
     source_t = torch.from_numpy(source.transpose(2, 0, 1)[None, ...]).to(device=device, dtype=torch.float32)
     source_edge = analysis.edge_map.astype(np.float32)
     edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge)
-    guide_edge = np.maximum(source_edge, _cluster_boundary_map(analysis.cluster_map))
-    edge_t = torch.from_numpy(guide_edge[None, None, ...]).to(device=device, dtype=torch.float32)
+    edge_t = torch.from_numpy(source_edge[None, None, ...]).to(device=device, dtype=torch.float32)
     uv0 = _make_regular_uv(
         height=height,
         width=width,

src/repixelizer/pipeline.py

@@ -7,7 +7,7 @@ from typing import Any
 import numpy as np
 
 from .analysis import analyze_continuous_source, analyze_tile_graph_source
-from .continuous import optimize_uv_field
+from .continuous import optimize_lattice_pixels
 from .diagnostics import (
     summarize_run,
     write_alpha_preview,
@@ -421,7 +421,7 @@ def _run_reconstruction(
             device=device,
             solver_params=solver_params,
         )
-    solver = optimize_uv_field(
+    solver = optimize_lattice_pixels(
         source,
         inference=inference,
         analysis=analysis,

src/repixelizer/types.py

@@ -30,7 +30,6 @@ class InferenceResult:
 @dataclass(slots=True)
 class ContinuousSourceAnalysis:
     edge_map: np.ndarray
-    cluster_map: np.ndarray
 
 
 @dataclass(slots=True)