Update src/streamlit_app.py

src/streamlit_app.py

@@ -55,32 +55,35 @@ class DominantColorDetector:
         self.edge_multiplier = edge_multiplier
 
     def _preprocess(self, img: Image.Image):
-        """Resize, keep aspect ratio, return float32 pixels."""
         if img.mode != "RGB":
             img = img.convert("RGB")
-        w, h = img.size
-        aspect = h / w
-        new_h = max(1, int(self.resize_dim * aspect))
-        img = img.resize((self.resize_dim, new_h), Image.LANCZOS)
-        return np.array(img).reshape(-1, 3).astype(np.float32)
-
-    def _build_edge_mask(self, img_array_flat, orig_shape):
-        """Create a weight mask highlighting edge regions."""
-        h, w, _ = orig_shape
-        mask = np.ones(h * w, dtype=np.float32)
-
-        # 10% strips at top/bottom/left/right
+            original_w, original_h = img.size
+            aspect = original_h / original_w
+            new_h = max(1, int(self.resize_dim * aspect))
+            img = img.resize((self.resize_dim, new_h), Image.LANCZOS)
+            resized_w, resized_h = img.size
+            pixels = np.array(img).reshape(-1, 3).astype(np.float32)
+            return pixels, (resized_w, resized_h)
+
+    def _build_edge_mask(self, total_pixels, resized_dims):
+        w, h = resized_dims
+        mask = np.ones(total_pixels, dtype=np.float32)
+        
         edge_frac = 0.10
         top_height = int(h * edge_frac)
         bottom_height = int(h * edge_frac)
         left_width = int(w * edge_frac)
         right_width = int(w * edge_frac)
-
+        
         for y in range(h):
             for x in range(w):
                 idx = y * w + x
-                if y < top_height or y >= h - bottom_height or \
-                   x < left_width or x >= w - right_width:
+                if (
+                    y < top_height
+                    or y >= h - bottom_height
+                    or x < left_width
+                    or x >= w - right_width
+                ):
                     mask[idx] = self.edge_multiplier
         return mask
 
@@ -181,75 +184,74 @@ class DominantColorDetector:
         }
 
     def detect_properties(self, img: Image.Image, include_palette=True):
-        orig_shape = img.size      # (width, height) before flatten
-        pixels_rgb = self._preprocess(img)
-        resized_shape = img.size   # after resize
-
-        # Build edge mask and create weighted sample pool
-        edge_mask = self._build_edge_mask(pixels_rgb, resized_shape[::-1])  # h,w
-        max_pixels = 10000
-        if len(pixels_rgb) > max_pixels:
-            # Weighted random choice
-            probs = edge_mask / edge_mask.sum()
-            idx = np.random.choice(len(pixels_rgb), max_pixels, replace=False, p=probs)
-            sampled_pixels = pixels_rgb[idx]
-            # Track edge influence for each sampled pixel
-            sampled_mask = edge_mask[idx]  # >1 for edges
-        else:
-            sampled_pixels = pixels_rgb
-            sampled_mask = edge_mask
-
-        pixels_lab = self._pixels_to_lab(sampled_pixels)
-        kmeans = KMeans(n_clusters=self.num_colors, random_state=42,
-                        n_init=3, max_iter=100, algorithm="elkan")
-        kmeans.fit(pixels_lab)
-        centroids_lab = kmeans.cluster_centers_
-        labels = kmeans.labels_
-        label_counts = Counter(labels)
-        total = len(labels)
-        centroids_rgb = self._lab_to_rgb(centroids_lab)
-
-        color_list = []
-        for i in range(self.num_colors):
-            lab = centroids_lab[i]
-            rgb = centroids_rgb[i]
-            pf = label_counts[i] / total
-            if pf < self.MIN_PIXEL_FRACTION:
-                continue
-            # Edge influence: fraction of this cluster's samples that were edge pixels
-            cluster_mask = labels == i
-            edge_frac = np.mean(sampled_mask[cluster_mask] > 1.0)  # proportion from edges
-            r, g, b = int(rgb[0]), int(rgb[1]), int(rgb[2])
-            color_list.append({
-                "lab": lab,
-                "rgb": {"red": r, "green": g, "blue": b},
-                "color": rgb_to_hex((r, g, b)),
-                "pixelFraction": float(pf),
-                "score": self._score_color(lab, pf),
-                "chroma": round(lab_chroma(lab), 2),
-                "isNearBlack": bool(self._is_near_black(lab)),
-                "isNearWhite": bool(self._is_near_white(lab)),
-                "isNearGray": bool(self._is_near_gray(lab)),
-                "edgeInfluence": round(edge_frac, 3),
-            })
-        color_list.sort(key=lambda x: x["score"], reverse=True)
-        color_list = self._remove_near_duplicates(color_list)
-        dominant_colors = [
-            {k: v for k, v in c.items() if k != "lab"}
-            for c in color_list
-        ]
-        for c in dominant_colors:
-            c["score"] = round(c["score"], 4)
-            c["pixelFraction"] = round(c["pixelFraction"], 4)
-
-        result = {
-            "imagePropertiesAnnotation": {
-                "dominantColors": {"colors": dominant_colors}
-            }
+    pixels_rgb, (resized_w, resized_h) = self._preprocess(img)
+    total_pixels = len(pixels_rgb)
+
+    # Build edge mask for the full resized image
+    edge_mask = self._build_edge_mask(total_pixels, (resized_w, resized_h))
+
+    max_pixels = 10000
+    if total_pixels > max_pixels:
+        probs = edge_mask / edge_mask.sum()
+        idx = np.random.choice(total_pixels, max_pixels, replace=False, p=probs)
+        sampled_pixels = pixels_rgb[idx]
+        sampled_mask = edge_mask[idx]
+    else:
+        sampled_pixels = pixels_rgb
+        sampled_mask = edge_mask
+
+    pixels_lab = self._pixels_to_lab(sampled_pixels)
+    kmeans = KMeans(n_clusters=self.num_colors, random_state=42,
+                    n_init=3, max_iter=100, algorithm="elkan")
+    kmeans.fit(pixels_lab)
+    centroids_lab = kmeans.cluster_centers_
+    labels = kmeans.labels_
+    label_counts = Counter(labels)
+    total = len(labels)
+    centroids_rgb = self._lab_to_rgb(centroids_lab)
+
+    color_list = []
+    for i in range(self.num_colors):
+        lab = centroids_lab[i]
+        rgb = centroids_rgb[i]
+        pf = label_counts[i] / total
+        if pf < self.MIN_PIXEL_FRACTION:
+            continue
+        cluster_mask = labels == i
+        edge_frac = np.mean(sampled_mask[cluster_mask] > 1.0)
+        r, g, b = int(rgb[0]), int(rgb[1]), int(rgb[2])
+        color_list.append({
+            "lab": lab,
+            "rgb": {"red": r, "green": g, "blue": b},
+            "color": rgb_to_hex((r, g, b)),
+            "pixelFraction": float(pf),
+            "score": self._score_color(lab, pf),
+            "chroma": round(lab_chroma(lab), 2),
+            "isNearBlack": bool(self._is_near_black(lab)),
+            "isNearWhite": bool(self._is_near_white(lab)),
+            "isNearGray": bool(self._is_near_gray(lab)),
+            "edgeInfluence": round(edge_frac, 3),
+        })
+
+    color_list.sort(key=lambda x: x["score"], reverse=True)
+    color_list = self._remove_near_duplicates(color_list)
+
+    dominant_colors = [
+        {k: v for k, v in c.items() if k != "lab"}
+        for c in color_list
+    ]
+    for c in dominant_colors:
+        c["score"] = round(c["score"], 4)
+        c["pixelFraction"] = round(c["pixelFraction"], 4)
+
+    result = {
+        "imagePropertiesAnnotation": {
+            "dominantColors": {"colors": dominant_colors}
         }
-        if include_palette:
-            result["suggestedPalette"] = self._build_adaptive_palette(color_list)
-        return result
+    }
+    if include_palette:
+        result["suggestedPalette"] = self._build_adaptive_palette(color_list)
+    return result
 
 # -------------------------------------------------------------------
 # Helper functions (global)