Update

README.md

@@ -4,7 +4,7 @@ emoji: 🏃
 colorFrom: gray
 colorTo: purple
 sdk: gradio
-sdk_version: 3.36.1
+sdk_version: 3.37.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -47,8 +47,9 @@ model = load_model()
 labels = load_labels()
 
 
-def predict(image: PIL.Image.Image,
-            score_threshold: float) -> dict[str, float]:
+def predict(
+        image: PIL.Image.Image, score_threshold: float
+) -> tuple[dict[str, float], dict[str, float], str]:
     _, height, width, _ = model.input_shape
     image = np.asarray(image)
     image = tf.image.resize(image,
@@ -60,12 +61,19 @@ def predict(image: PIL.Image.Image,
     image = image / 255.
     probs = model.predict(image[None, ...])[0]
     probs = probs.astype(float)
-    res = dict()
-    for prob, label in zip(probs.tolist(), labels):
+
+    indices = np.argsort(probs)[::-1]
+    result_all = dict()
+    result_threshold = dict()
+    for index in indices:
+        label = labels[index]
+        prob = probs[index]
+        result_all[label] = prob
         if prob < score_threshold:
-            continue
-        res[label] = prob
-    return res
+            break
+        result_threshold[label] = prob
+    result_text = ', '.join(result_all.keys())
+    return result_threshold, result_all, result_text
 
 
 image_paths = load_sample_image_paths()
@@ -83,15 +91,26 @@ with gr.Blocks(css='style.css') as demo:
                                         value=0.5)
             run_button = gr.Button('Run')
         with gr.Column():
-            result = gr.Label(label='Output')
+            with gr.Tabs():
+                with gr.Tab(label='Output'):
+                    result = gr.Label(label='Output', show_label=False)
+                with gr.Tab(label='JSON'):
+                    result_json = gr.JSON(label='JSON output',
+                                          show_label=False)
+                with gr.Tab(label='Text'):
+                    result_text = gr.Text(label='Text output',
+                                          show_label=False,
+                                          lines=5)
     gr.Examples(examples=examples,
                 inputs=[image, score_threshold],
-                outputs=result,
+                outputs=[result, result_json, result_text],
                 fn=predict,
                 cache_examples=os.getenv('CACHE_EXAMPLES') == '1')
 
-    run_button.click(fn=predict,
-                     inputs=[image, score_threshold],
-                     outputs=result,
-                     api_name='predict')
+    run_button.click(
+        fn=predict,
+        inputs=[image, score_threshold],
+        outputs=[result, result_json, result_text],
+        api_name='predict',
+    )
 demo.queue().launch()

requirements.txt

@@ -1,3 +1,4 @@
-pillow>=9.0.0
-tensorflow>=2.7.0
 git+https://github.com/KichangKim/DeepDanbooru@v3-20200915-sgd-e30#egg=deepdanbooru
+pillow==10.0.0
+pydantic==1.10.11
+tensorflow==2.13.0