dev(narugo): add text export

app.py

@@ -1,5 +1,6 @@
+import re
 from functools import lru_cache
-from typing import List, Mapping
+from typing import List, Mapping, Tuple
 
 import gradio as gr
 import numpy as np
@@ -48,27 +49,55 @@ def image_preprocess(image: Image.Image) -> np.ndarray:
     return data.reshape((1, 512, 512, 3))  # B x H x W x C
 
 
-def image_to_deepdanbooru_tags(image: Image.Image, threshold: float) -> Mapping[str, float]:
+RE_SPECIAL = re.compile(r'([\\()])')
+
+
+def image_to_deepdanbooru_tags(image: Image.Image, threshold: float,
+                               use_spaces: bool, use_escape: bool, include_ranks: bool, score_descend: bool) \
+        -> Tuple[Mapping[str, float], str]:
     tags = get_deepdanbooru_tags()
     session = get_deepdanbooru_onnx()
     input_name = session.get_inputs()[0].name
     output_names = [output.name for output in session.get_outputs()]
 
     result = session.run(output_names, {input_name: image_preprocess(image)})[0]
-    return {
+    filtered_tags = {
         tag: float(score) for tag, score in zip(tags, result[0])
         if score >= threshold
     }
 
+    text_items = []
+    tags_pairs = filtered_tags.items()
+    if score_descend:
+        tags_pairs = sorted(tags_pairs, key=lambda x: (-x[1], x[0]))
+    for tag, score in tags_pairs:
+        tag_outformat = tag
+        if use_spaces:
+            tag_outformat = tag_outformat.replace('_', ' ')
+        if use_escape:
+            tag_outformat = re.sub(RE_SPECIAL, r'\\\1', tag_outformat)
+        if include_ranks:
+            tag_outformat = f"({tag_outformat}:{score:.3f})"
+        text_items.append(tag_outformat)
+    output_text = ', '.join(text_items)
+
+    return output_text, filtered_tags
+
 
 if __name__ == '__main__':
     interface = gr.Interface(
         image_to_deepdanbooru_tags,
         inputs=[
             gr.Image(type='pil', label='Original Image'),
-            gr.Slider(0.0, 1.0, 0.5, label='Tagging Confidence Threshold')],
+            gr.Slider(0.0, 1.0, 0.5, label='Tagging Confidence Threshold'),
+            gr.Checkbox(value=False, label='Use Space Instead Of _'),
+            gr.Checkbox(value=True, label='Use Text Escape'),
+            gr.Checkbox(value=False, label='Keep Confidences'),
+            gr.Checkbox(value=True, label='Descend By Confidence'),
+        ],
         outputs=[
-            gr.Label(),
+            gr.TextArea(label='Exported Text'),
+            gr.Label(label='Tags'),
         ],
         interpretation="default"