data

app.py

@@ -1,5 +1,27 @@
 import gradio as gr
 import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class RangeModel(nn.Module):
+    def __init__(self):
+        super(RangeModel, self).__init__()
+        self.cross_attn = Attn(128, 8192)
+        self.upsample = nn.Upsample(scale_factor=4, mode='bilinear')
+        self.out = nn.Conv2d(128, 1, 1, bias=False)
+        self.x = None
+    
+    def forward(self, text):
+        x = self.cross_attn(self.x, text)
+        x = rearrange(x, 'b (h w) d -> b d h w', h=225)
+        x = self.upsample(x)
+        x = self.out(x)
+        return x
+
+model = RangeModel()
+model.load_state_dict(torch.load("model/demo_model.pt", map_location=torch.device('cpu')))
 
 species = np.load("data/species_70b.npy", allow_pickle=True)
 clas = np.load("data/class_70b.npy", allow_pickle=True)
@@ -13,7 +35,7 @@ order_list = list(order[()].keys())
 #genus_list = list(genus[()].keys())
 #family_list = list(family[()].keys())
 
-pos_embed = np.load("data/pos_embed.npy", allow_pickle=True)
+#pos_embed = np.load("data/pos_embed.npy", allow_pickle=True)
 
 def update_fn(val):
     if val=="Class":
@@ -38,7 +60,10 @@ def pred_fn(taxon, name):
         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)
 
 
 with gr.Blocks() as demo:
@@ -60,5 +85,4 @@ with gr.Blocks() as demo:
     
     submit_button.click(pred_fn, inputs=[inp, out], outputs=[pred])
     
-
 demo.launch()