Create app.py

app.py

@@ -0,0 +1,85 @@
+import io
+from pprint import pformat
+
+import gradio as gr
+from hbutils.string import titleize
+from hfutils.repository import hf_hub_repo_url
+
+from imgutils.tagging.pixai import _open_default_category_thresholds, get_pixai_tags
+
+REPO_ID = 'deepghs/pixai-tagger-v0.9-onnx'
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                repo_url = hf_hub_repo_url(repo_id=REPO_ID, repo_type='model')
+                gr.HTML(f'<h2 style="text-align: center;">Tagger Demo For {REPO_ID}</h2>')
+                gr.Markdown(f'This is the quick demo for tagger model [{REPO_ID}]({repo_url}). '
+                            f'Powered by `dghs-imgutils`\'s quick demo module.')
+
+        with gr.Row():
+            thresholds, names = _open_default_category_thresholds(model_name=REPO_ID)
+            categories = sorted(set(names.keys()))
+
+            with gr.Column():
+                with gr.Row():
+                    gr_input_image = gr.Image(type='pil', label='Original Image')
+                with gr.Row():
+                    gr_thresholds = []
+                    for category in categories:
+                        gr_cate_threshold = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=thresholds[category],
+                            step=0.001,
+                            label=f'Threshold for {titleize(names[category])}',
+                        )
+                        gr_thresholds.append(gr_cate_threshold)
+                with gr.Row():
+                    gr_submit = gr.Button(value='Submit', variant='primary')
+
+            with gr.Column():
+                with gr.Tabs():
+                    gr_preds = []
+                    for category in categories:
+                        with gr.Tab(f'{titleize(names[category])}'):
+                            gr_cate_label = gr.Label(f'{titleize(names[category])} Prediction')
+                            gr_preds.append(gr_cate_label)
+                    with gr.Tab('IPs Mapping'):
+                        gr_ips_mapping = gr.TextArea(label="IPs (string)", lines=15)
+                    with gr.Tab('Text Output'):
+                        gr_text_output = gr.TextArea(label="Output (string)", lines=15)
+
+
+            def _fn_submit(image, *thresholds):
+                _ths = {
+                    category: cate_ths
+                    for category, cate_ths in zip(categories, thresholds)
+                }
+
+                fmt = {
+                    **names,
+                    'ips_mapping': 'ips_mapping',
+                    'ips': 'ips',
+                }
+                res = get_pixai_tags(image=image, model_name=REPO_ID, thresholds=_ths, fmt=fmt)
+                with io.StringIO() as sf:
+                    for category in categories:
+                        print(f'# {names[category]} (#{category})', file=sf)
+                        print(f', '.join(res[category].keys()), file=sf)
+                        print(f'', file=sf)
+                    print(f'# IPs', file=sf)
+                    print(f', '.join(res['ips']), file=sf)
+                    print(f'', file=sf)
+                    return sf.getvalue(), pformat(res['ips_mapping']), \
+                        *[res[category] for category in categories]
+
+
+            gr_submit.click(
+                fn=_fn_submit,
+                inputs=[gr_input_image, *gr_thresholds],
+                outputs=[gr_text_output, gr_ips_mapping, *gr_preds]
+            )
+
+    demo.launch()