Return a list of predicted labels

app.py

@@ -7,6 +7,7 @@ import functools
 import os
 import pathlib
 import tarfile
+import tempfile
 
 import deepdanbooru as dd
 import gradio as gr
@@ -64,7 +65,8 @@ def load_labels() -> list[str]:
 
 
 def predict(image: PIL.Image.Image, score_threshold: float,
-            model: tf.keras.Model, labels: list[str]) -> dict[str, float]:
+            model: tf.keras.Model,
+            labels: list[str]) -> tuple[dict[str, float], str]:
     _, height, width, _ = model.input_shape
     image = np.asarray(image)
     image = tf.image.resize(image,
@@ -81,7 +83,14 @@ def predict(image: PIL.Image.Image, score_threshold: float,
         if prob < score_threshold:
             continue
         res[label] = prob
-    return res
+
+    sorted_preds = sorted(res.items(), key=lambda x: -x[1])
+    out_path = tempfile.NamedTemporaryFile(suffix='.txt', delete=False)
+    with open(out_path.name, 'w') as f:
+        for key, _ in sorted_preds:
+            f.write(f'{key}\n')
+
+    return res, out_path.name
 
 
 def main():
@@ -106,7 +115,10 @@ def main():
                       value=args.score_threshold,
                       label='Score Threshold'),
         ],
-        gr.Label(label='Output'),
+        [
+            gr.Label(label='Output'),
+            gr.File(label='Tag List'),
+        ],
         examples=examples,
         title=TITLE,
         description=DESCRIPTION,