Spaces:
Sleeping
Sleeping
File size: 2,571 Bytes
67a423d dae87e0 67a423d dae87e0 67a423d 23a411c 67a423d dae87e0 06cef48 67a423d dae87e0 06cef48 dae87e0 67a423d dae87e0 67a423d dae87e0 67a423d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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()
|