dev(narugo): save this

app.py

@@ -2,6 +2,7 @@ import os
 
 import gradio as gr
 import numpy as np
+import pandas as pd
 from PIL import Image
 from sklearn.cluster import KMeans
 
@@ -27,10 +28,12 @@ def get_main_colors(image: Image.Image, n: int = 28, pixels: int = 90000) -> Ima
 
     width, height = image.size
     raw = np.asarray(image).reshape(-1, 3)
-    new_data = kmeans.cluster_centers_[kmeans.predict(raw)]
-    new_data = new_data.round().astype(np.uint8).reshape((height, width, 3))
+    colors = kmeans.cluster_centers_.round().astype(np.uint8)
+    prediction = kmeans.predict(raw)
+    new_data = colors[prediction].reshape((height, width, 3))
+    new_image = Image.fromarray(new_data, mode='RGB')
 
-    return Image.fromarray(new_data, mode='RGB')
+    return new_image
 
 
 def main_func(image: Image.Image, n: int, pixels: int, fixed_width: bool, width: int):
@@ -41,10 +44,29 @@ def main_func(image: Image.Image, n: int, pixels: int, fixed_width: bool, 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)
 
-    return new_image
+    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['ratio'] = table['count'] / table['count'].sum()
+    hexes = []
+    for r, g, b in zip(table['r'], table['g'], table['b']):
+        hexes.append(f'#{r:02x}{g:02x}{b:02x}')
+    table['hex'] = hexes
+
+    new_table = pd.DataFrame({
+        'Hex': table['hex'],
+        'Pixels': table['count'],
+        'Ratio': table['ratio'],
+        'Red': table['r'],
+        'Green': table['g'],
+        'Blue': table['b'],
+    })
+
+    return new_image, new_table
 
 
 if __name__ == '__main__':
+    pd.set_option("display.precision", 3)
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
@@ -59,12 +81,19 @@ if __name__ == '__main__':
                 ch_submit = gr.Button(value='Submit', variant='primary')
 
             with gr.Column():
-                ch_output = gr.Image(type='pil', label='Output Image')
+                with gr.Tabs():
+                    with gr.Tab('Output Image'):
+                        ch_output = gr.Image(type='pil', label='Output Image')
+                    with gr.Tab('Color Map'):
+                        ch_color_map = gr.Dataframe(
+                            headers=['Hex', 'Pixels', 'Ratio', 'Red', 'Green', 'Blue'],
+                            label='Color Map'
+                        )
 
         ch_submit.click(
             main_func,
             inputs=[ch_image, ch_clusters, ch_pixels, ch_fixed_width, ch_width],
-            outputs=[ch_output],
+            outputs=[ch_output, ch_color_map],
         )
 
     demo.queue(os.cpu_count()).launch()