data

.gitattributes

@@ -39,3 +39,4 @@ data/species_70b.npy filter=lfs diff=lfs merge=lfs -text
 data/pos_embeds_model.npy filter=lfs diff=lfs merge=lfs -text
 model/demo_model.pt filter=lfs diff=lfs merge=lfs -text
 data/family_70b.npy filter=lfs diff=lfs merge=lfs -text
+data/LAND_MASK.tif filter=lfs diff=lfs merge=lfs -text

app.py

@@ -5,9 +5,13 @@ import torch.nn as nn
 from einops import rearrange
 import matplotlib.pyplot as plt
 from copy import deepcopy
+import rasterio
+from rasterio.enums import Resampling
 
 pred_global = None
 
+land_mask = (rasterio.open('data/LAND_MASK.tif').read(out_shape=(1, 900, 1800), resampling=Resampling.nearest) == 1).squeeze(0)
+
 class Attn(nn.Module):
     def __init__(self, dim, dim_text, heads = 16, dim_head = 64):
         super().__init__()
@@ -94,7 +98,7 @@ def text_fn(taxon, name):
         text_embeds = species[()][name]
     
     text_embeds = torch.tensor(text_embeds)
-    preds = model(text_embeds).sigmoid().squeeze(0).squeeze(0).detach().numpy()
+    preds = model(text_embeds).sigmoid().squeeze(0).squeeze(0).detach().numpy() * land_mask
     pred_global = preds
     cmap = plt.get_cmap('Greens')
 

requirements.txt

@@ -2,4 +2,5 @@ numpy==1.23.4
 torch==2.0.1
 rasterio==1.3.8
 einops==0.6.1
-matplotlib
+matplotlib
+rasterio