Fixes

.gitignore

@@ -1,2 +1,4 @@
 .venv
-flagged
+flagged
+*.tif
+*.tiff

app.py

@@ -1,12 +1,11 @@
-from io import BytesIO
 import os
+import pickle
 import re
 import PIL.Image
 import pandas as pd
 import numpy as np
 import gradio as gr
 from datasets import load_dataset
-import infer
 import matplotlib.pyplot as plt
 from sklearn.manifold import TSNE
 from sklearn.preprocessing import LabelEncoder
@@ -15,8 +14,10 @@ from torch import nn
 from transformers import BertConfig, BertForMaskedLM, PreTrainedTokenizerFast
 from huggingface_hub import PyTorchModelHubMixin
 from pinecone import Pinecone
+import rasterio
+from rasterio.sample import sample_gen
 
-from config import DEFAULT_INPUTS, MODELS, DATASETS, ID_TO_GENUS_MAP
+from config import DEFAULT_INPUTS, MODELS, DATASETS, ID_TO_GENUS_MAP, LAYER_NAMES
 
 # We need this for the eco layers because they are too big
 PIL.Image.MAX_IMAGE_PIXELS = None
@@ -52,15 +53,15 @@ classification_model = DNASeqClassifier.from_pretrained(
         BertConfig(vocab_size=259, output_hidden_states=True),
     ),
 )
+with open("scaler.pkl", "rb") as f:
+    scaler = pickle.load(f)
 
 embeddings_model.eval()
 classification_model.eval()
 
 # Load datasets
-ecolayers_ds = load_dataset(DATASETS["ecolayers"])
 amazon_ds = load_dataset(DATASETS["amazon"])
 
-
 def set_default_inputs():
     return (DEFAULT_INPUTS["dna_sequence"],
             DEFAULT_INPUTS["latitude"],
@@ -99,7 +100,6 @@ def tokenize(dna_sequence: str) -> dict[str, torch.Tensor]:
     return tokenizer(dna_seq_preprocessed, return_tensors="pt")
 
 
-
 def get_embedding(dna_sequence: str) -> torch.Tensor:
     dna_embedding: torch.Tensor = embeddings_model(
         **tokenize(dna_sequence)
@@ -126,7 +126,20 @@ def predict_genus(method: str, dna_sequence: str, latitude: str, longitude: str)
     
     if method == "fine_tuned_model":
         bert_inputs = tokenize(dna_sequence)
-        logits = classification_model(bert_inputs, torch.zeros(1, 7))
+
+        env_data = []
+        for layer in LAYER_NAMES:
+            with rasterio.open(layer) as dataset:
+                # Get the corresponding ecological values for the samples
+                results = sample_gen(dataset, [coords])
+                results = [r for r in results]
+            layer_data = np.mean(results[0])
+            env_data.append(layer_data)
+
+        env_data = scaler.transform([env_data])
+        env_data = torch.from_numpy(env_data).to(torch.float32)
+
+        logits = classification_model(bert_inputs, env_data)
         temperature = 0.2
         probs = torch.softmax(logits / temperature, dim=1).squeeze()
         top_k = torch.topk(probs, 10)

config.py

@@ -1,3 +1,4 @@
+import os
 import json
 
 
@@ -24,6 +25,17 @@ MODELS = {
 }
 
 DATASETS = {
-    "ecolayers": "LofiAmazon/Global-Ecolayers",
     "amazon": "LofiAmazon/BOLD-Embeddings-Ecolayers-Amazon",
 }
+
+HUGGINGFACE_DW_URL = "https://huggingface.co/datasets/LofiAmazon/Global-Ecolayers/resolve/main/{filename}?download=true"
+
+LAYER_NAMES = [
+    "median_elevation_1km.tiff",
+    "human_footprint.tiff",
+    "population_density_1km.tif",
+    "annual_precipitation.tif",
+    "precipitation_seasonality.tif",
+    "annual_mean_air_temp.tif",
+    "temp_seasonality.tif",
+]