dev(hansbug): update code

app.py

@@ -1,4 +1,5 @@
 import os
+from typing import Tuple
 
 import gradio as gr
 import numpy as np
@@ -16,7 +17,8 @@ def _image_resize(image: Image.Image, pixels: int = 90000, **kwargs):
     return small_image
 
 
-def get_main_colors(image: Image.Image, n: int = 28, pixels: int = 90000) -> Image.Image:
+def get_main_colors(image: Image.Image, n: int = 28, pixels: int = 90000) \
+        -> Tuple[Image.Image, np.ndarray, np.ndarray, np.ndarray]:
     image = image.copy()
     if image.mode != 'RGB':
         image = image.convert('RGB')
@@ -33,20 +35,27 @@ def get_main_colors(image: Image.Image, n: int = 28, pixels: int = 90000) -> Ima
     new_data = colors[prediction].reshape((height, width, 3))
     new_image = Image.fromarray(new_data, mode='RGB')
 
-    return new_image
+    cids, counts = np.unique(prediction, return_counts=True)
+    counts = np.asarray(list(map(lambda x: x[1], sorted(zip(cids, counts)))))
+
+    return new_image, colors, counts, prediction.reshape((height, width))
 
 
 def main_func(image: Image.Image, n: int, pixels: int, fixed_width: bool, width: int):
-    new_image = get_main_colors(image, n, pixels)
     if fixed_width:
-        _width, _height = new_image.size
+        _width, _height = image.size
         r = width / _width
         new_width, new_height = int(round(_width * r)), int(round(_height * r))
-        new_image = new_image.resize((new_width, new_height), resample=Image.NEAREST)
+        image = image.resize((new_width, new_height))
+
+    new_image, colors, counts, predictions = get_main_colors(image, n, pixels)
 
-    df = pd.DataFrame(data=np.asarray(new_image).reshape(-1, 3), columns=['r', 'g', 'b'])
-    df['id_'] = 1
-    table = df.groupby(['r', 'g', 'b'])['id_'].agg(['count']).reset_index().sort_values('count', ascending=False)
+    table = pd.DataFrame({
+        'r': colors[:, 0],
+        'g': colors[:, 1],
+        'b': colors[:, 2],
+        'count': counts,
+    })
     table['ratio'] = table['count'] / table['count'].sum()
     hexes = []
     for r, g, b in zip(table['r'], table['g'], table['b']):
@@ -60,7 +69,7 @@ def main_func(image: Image.Image, n: int, pixels: int, fixed_width: bool, width:
         'Red': table['r'],
         'Green': table['g'],
         'Blue': table['b'],
-    })
+    }).sort_values('Pixels', ascending=False)
 
     return new_image, new_table