Upload anime_aesthetic/app.py with huggingface_hub

anime_aesthetic/app.py

@@ -0,0 +1,84 @@
+import os
+from functools import lru_cache
+
+import gradio as gr
+import numpy as np
+import pandas as pd
+from huggingface_hub import HfFileSystem, hf_hub_download
+from imgutils.generic import classify_predict_score
+from natsort import natsorted
+
+hf_fs = HfFileSystem()
+
+_REPOSITORY = 'deepghs/anime_aesthetic'
+_DEFAULT_MODEL = 'swinv2pv3_v0_448_ls0.2_x'
+_MODELS = natsorted([
+    os.path.dirname(os.path.relpath(file, _REPOSITORY))
+    for file in hf_fs.glob(f'{_REPOSITORY}/*/model.onnx')
+])
+
+LABELS = ["worst", "low", "normal", "good", "great", "best", "masterpiece"]
+
+
+@lru_cache()
+def _get_mark_table(model):
+    df = pd.read_csv(hf_hub_download(
+        repo_id=_REPOSITORY,
+        repo_type='model',
+        filename=f'{model}/samples.csv',
+    ))
+    df = df.sort_values(['score'])
+    df['cnt'] = list(range(len(df)))
+    df['final_score'] = df['cnt'] / len(df)
+
+    x = np.concatenate([[0.0], df['score'], [6.0]])
+    y = np.concatenate([[0.0], df['final_score'], [1.0]])
+    return x, y
+
+
+def _get_percentile(x, y, v):
+    idx = np.searchsorted(x, np.clip(v, a_min=0.0, a_max=6.0))
+    if idx < x.shape[0] - 1:
+        x0, y0 = x[idx], y[idx]
+        x1, y1 = x[idx + 1], y[idx + 1]
+        return np.clip((v - x0) / (x1 - x0) * (y1 - y0) + y0, a_min=0.0, a_max=1.0)
+
+    else:
+        return y[idx]
+
+
+def _fn_predict(image, model):
+    scores = classify_predict_score(
+        image=image,
+        repo_id=_REPOSITORY,
+        model_name=model,
+    )
+    weighted_mean = sum(i * scores[label] for i, label in enumerate(LABELS))
+    x, y = _get_mark_table(model)
+    percentile = _get_percentile(x, y, weighted_mean)
+    return weighted_mean, percentile, scores
+
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                gr_input_image = gr.Image(type='pil', label='Original Image')
+                gr_model = gr.Dropdown(_MODELS, value=_DEFAULT_MODEL, label='Model')
+                gr_submit = gr.Button(value='Submit', variant='primary')
+
+            with gr.Column():
+                with gr.Row():
+                    gr_score = gr.Text(label='Aesthetic Score (0~6)', value='')
+                    gr_percentile = gr.Text(label='Percentile (0.0-1.0)', value='')
+
+                with gr.Row():
+                    gr_output = gr.Label(label='Aesthetic Classes')
+
+            gr_submit.click(
+                _fn_predict,
+                inputs=[gr_input_image, gr_model],
+                outputs=[gr_score, gr_percentile, gr_output],
+            )
+
+    demo.queue(os.cpu_count()).launch()