dev(narugo): add text export feature on gradio

app.py

@@ -1,3 +1,4 @@
+import re
 from typing import Mapping, Tuple, Dict
 
 import cv2
@@ -118,17 +119,37 @@ WAIFU_MODELS: Mapping[str, WaifuDiffusionInterrogator] = {
         repo='SmilingWolf/wd-v1-4-convnext-tagger'
     ),
 }
+RE_SPECIAL = re.compile(r'([\\()])')
 
 
-def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float):
+def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
+                       use_spaces: bool, use_escape: bool, include_ranks: bool, score_descend: bool) \
+        -> Tuple[Mapping[str, float], str, Mapping[str, float]]:
     model = WAIFU_MODELS[model_name]
     ratings, tags = model.interrogate(image)
 
-    return ratings, {
+    filtered_tags = {
         tag: score for tag, score in tags.items()
         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 ratings, output_text, filtered_tags
+
 
 if __name__ == '__main__':
     interface = gr.Interface(
@@ -136,10 +157,15 @@ if __name__ == '__main__':
         inputs=[
             gr.Image(type='pil', label='Original Image'),
             gr.Radio(list(WAIFU_MODELS.keys()), value='wd14-vit', label='Waifu Model'),
-            gr.Slider(0.0, 1.0, value=0.5, label='Tagging Confidence Threshold')
+            gr.Slider(0.0, 1.0, value=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(label='Ratings'),
+            gr.TextArea(label='Exported Text'),
             gr.Label(label='Tags'),
         ],
         interpretation="default"