data

app.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 from einops import rearrange
 import matplotlib.pyplot as plt
 
+pred_global = None
 
 class Attn(nn.Module):
     def __init__(self, dim, dim_text, heads = 16, dim_head = 64):
@@ -79,6 +80,7 @@ def update_fn(val):
         return gr.Dropdown(label="Name", choices=species_list, interactive=True)
 
 def text_fn(taxon, name):
+    global pred_global
     if taxon=="Class":
         text_embeds = clas[()][name]
     elif taxon=="Order":
@@ -92,6 +94,7 @@ def text_fn(taxon, name):
     
     text_embeds = torch.tensor(text_embeds)
     preds = model(text_embeds).sigmoid().squeeze(0).squeeze(0).detach().numpy()
+    pred_global = preds
     cmap = plt.get_cmap('Greens')
 
     rgba_img = cmap(preds)
@@ -99,22 +102,14 @@ def text_fn(taxon, name):
     #return gr.Image(preds, label="Predicted Heatmap", visible=True)
     return rgb_img
 
-def pred_fn(taxon, name):
-    if taxon=="Class":
-        text_embeds = clas[()][name]
-    elif taxon=="Order":
-        text_embeds = order[()][name]
-    elif taxon=="Family":
-        text_embeds = family[()][name]
-    elif taxon=="Genus":
-        text_embeds = genus[()][name]
-    elif taxon=="Species":
-        text_embeds = species[()][name]
-    
-    text_embeds = torch.tensor(text_embeds)
-    preds = model(text_embeds).sigmoid().unsqueeze(0).unsqueeze(0).detach().numpy()
-    return gr.Image(preds, label="Predicted Heatmap", visible=True)
+def thresh_fn(val):
+    global pred_global
+    pred_global = pred_global > val
+    cmap = plt.get_cmap('Greens')
 
+    rgba_img = cmap(pred_global)
+    rgb_img = np.delete(rgba_img, 3, 2)
+    return rgb_img
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -137,5 +132,6 @@ with gr.Blocks() as demo:
         pred = gr.Image(label="Predicted Heatmap", visible=True)
     
     check_button.click(text_fn, inputs=[inp, out], outputs=[pred])
+    slider.change(thresh_fn, slider, pred)
     
 demo.launch()