major UI overhaul

app.py

@@ -128,23 +128,33 @@ allowed_tags = list(tags.keys())
 for idx, tag in enumerate(allowed_tags):
     allowed_tags[idx] = tag.replace("_", " ")
 
+sorted_tag_score = {}
+
 @spaces.GPU(duration=5)
-def create_tags(image, threshold):
+def run_classifier(image, threshold):
+    global sorted_tag_score
     img = image.convert('RGB')
     tensor = transform(img).unsqueeze(0)
 
     with torch.no_grad():
         logits = model(tensor)
         probabilities = torch.nn.functional.sigmoid(logits[0])
-        indices = torch.where(probabilities > threshold)[0]
+        indices = torch.where(probabilities, 250).indices
         values = probabilities[indices]
 
     tag_score = dict()
     for i in range(indices.size(0)):
         tag_score[allowed_tags[indices[i]]] = values[i].item()
     sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))
-    text_no_impl = ", ".join(sorted_tag_score.keys())
-    return text_no_impl, sorted_tag_score
+
+    return create_tags(threshold)
+
+def create_tags(threshold):
+    global sorted_tag_score
+    filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}
+    text_no_impl = ", ".join(filtered_tag_score.keys())
+    return text_no_impl, filtered_tag_score
+    
 
 with gr.Blocks(css=".output-class { display: none; }") as demo:
     gr.Markdown("""
@@ -153,14 +163,24 @@ with gr.Blocks(css=".output-class { display: none; }") as demo:
 
     This tagger is the result of joint efforts between members of the RedRocket team.  Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
     """)
-    gr.Interface(
-        create_tags,
-        inputs=[gr.Image(label="Source", sources=['upload'], type='pil', height="50vh", show_label=False), gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")],
-        outputs=[
-            gr.Textbox(label="Tag String"),
-            gr.Label(label="Tag Predictions", num_top_classes=200, show_label=False),
-        ],
-        allow_flagging="never",
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Source", sources=['upload'], type='pil', height="60vh", show_label=False)
+            threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+        with gr.Column():
+            tag_string = gr.Textbox(label="Tag String")
+            label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)
+
+    image_input.upload(
+        fn=run_classifier,
+        inputs=[image_input, threshold_slider],
+        outputs=[tag_string, label_box]
+    )
+
+    threshold_slider.input(
+        fn=create_tags,
+        inputs=[threshold_slider],
+        outputs=[tag_string, label_box]
     )
 
 if __name__ == "__main__":