Update app.py

app.py

@@ -58,6 +58,7 @@ tfm = transforms.Compose([
     ),
 ])
 
+
 @torch.no_grad()
 def predict(image: Image.Image):
     if image is None:
@@ -68,8 +69,8 @@ def predict(image: Image.Image):
 
     log_sp, log_st = model(x)
 
-    prob_sp = torch.softmax(log_sp, dim=1)[0]  # [num_species]
-    prob_st = torch.softmax(log_st, dim=1)[0]  # [num_states]
+    prob_sp = torch.softmax(log_sp, dim=1)[0]
+    prob_st = torch.softmax(log_st, dim=1)[0]
 
     sp_id = int(prob_sp.argmax().item())
     st_id = int(prob_st.argmax().item())
@@ -82,48 +83,65 @@ def predict(image: Image.Image):
 
     return sp_text, st_text
 
+
 EXAMPLE_TEXT = """
-12 taxa: Cacoxnus indagator, Chelostoma florisomne, Coeliopencyrtus, Eumenidae,
-Heriades, Hylaeus, Megachile, Osmia bicornis, Osmia cornuta, Passaloecus, Psenulus, Trypoxylon
+**Taxa (12):** Cacoxnus indagator, Chelostoma florisomne, Coeliopencyrtus, Eumenidae, Heriades, Hylaeus, Megachile, Osmia bicornis, Osmia cornuta, Passaloecus, Psenulus, Trypoxylon
 
-4 states: DauLv, DeadLv, Lv, OldFood
+**States (your model labels):** DauLv, DeadLv, Lv, OldFood
 """
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Textbox(label="Predicted taxon"),
-        gr.Textbox(label="Predicted state"),
-    ],
-    title="EfficientNet Two-Head Layer Trap Nest (LTN) Classifier",
-    description=EXAMPLE_TEXT,
-    examples=[
-        ["./0.jpg"],
-        ["./1.jpg"],
-        ["./2.jpg"],
-        ["./3.jpg"],
-        ["./4.jpg"],
-        ["./5.jpg"],
-        ["./6.jpg"],
-        ["./7.jpg"],
-        ["./8.jpg"],
-        ["./9.jpg"],
-        ["./10.jpg"],
-        ["./11.jpg"],
-        ["./12.jpg"],
-        ["./13.jpg"],
-        ["./14.jpg"],
-        ["./15.jpg"],
-        ["./16.jpg"],
-        ["./17.jpg"],
-        ["./18.jpg"],
-        ["./19.jpg"],
-        ["./20.jpg"],
-        ["./21.jpg"],
-        ["./22.jpg"],
-    ],
-)
+STATUS_TABLE = [
+    ["Dead",    "D", "DeadLv",  "Dead visible larva"],
+    ["Food",    "F", "OldFood", "Brood cell with only unconsumed food, no larva will develop"],
+    ["Hatched", "H", "Hatched", "Brood cell with hatched bee or wasp traces (cocoon or exuvia)"],
+    ["Larva",   "L", "Lv",      "Visible alive larva"],
+    ["Prepupa", "P", "Prepupa", "Visible alive prepupa that stopped feeding"],
+]
+
+
+theme = gr.themes.Soft()
+
+with gr.Blocks(theme=theme, title="LTN EfficientNet Two-Head Classifier") as demo:
+    gr.Markdown(
+        "# EfficientNet Two-Head Layer Trap Nest (LTN) Classifier\n"
+        "Upload a brood-cell image crop to predict **taxon** and **developmental status**.\n"
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            img = gr.Image(type="pil", label="Input image", height=320)
+
+            gr.Markdown("### Label sets")
+            gr.Markdown(EXAMPLE_TEXT)
+
+            gr.Markdown("### Status legend")
+            legend = gr.Dataframe(
+                value=STATUS_TABLE,
+                headers=["Status", "Letter", "Code", "Description"],
+                datatype=["str", "str", "str", "str"],
+                interactive=False,
+                wrap=True,
+                row_count=(len(STATUS_TABLE), "fixed"),
+                col_count=(4, "fixed"),
+            )
+
+        with gr.Column(scale=1):
+            sp_out = gr.Textbox(label="Predicted taxon", lines=1)
+            st_out = gr.Textbox(label="Predicted state", lines=1)
+
+            btn = gr.Button("Predict", variant="primary")
+            btn.click(fn=predict, inputs=img, outputs=[sp_out, st_out])
+
+            gr.Markdown("### Examples")
+            gr.Examples(
+                examples=[[f"./{i}.jpg"] for i in range(23)],
+                inputs=img,
+            )
+
+    gr.Markdown(
+        "<small>Note: The status legend is a human-readable mapping. "
+        "Your model output labels come from `labels.json`.</small>"
+    )
 
 if __name__ == "__main__":
     demo.launch()