Use Clay embeddings to predict earthquake based on Sentinel-1 imagery

[anindabitm]

Hi everyone,

Thanks for providing the wonderful Clay foundation model. I tried using the Clay embeddings for earthquake prediction using Sentinel-1 imagery. The dataset used is - "QuakeSet: A Dataset and Low-Resource Models to Monitor Earthquakes through Sentinel-1". I got decent F1-score using the Clay embeddings using a general Random Forest Classifier. My sense is we could get better results using xgbooost or full fine tune of the model. This goes to show the strength of Clay embeddings even for Sentinel-1 data. The notebook is attached below.

import torchgeo
import torchgeo.datasets as datasets
import sys
import torch
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt


sys.path.append("model/")

from src.model import ClayMAEModule

torchgeo.__version__

### Downloading the QuakeSet dataset

@misc{cambrin2024quakesetdatasetlowresourcemodels, title={QuakeSet: A Dataset and Low-Resource Models to Monitor Earthquakes through Sentinel-1}, author={Daniele Rege Cambrin and Paolo Garza}, year={2024}, eprint={2403.18116}, archivePrefix={arXiv}, primaryClass={cs.CV}, url={[https://arxiv.org/abs/2403.18116}](https://www.google.com/url?q=https%3A%2F%2Farxiv.org%2Fabs%2F2403.18116%257D), }

train_ds = datasets.QuakeSet(split="train",download=True)

#Checking the Sentinel-1 imagery size. 

#Each sample consists of pre and post event data with two channels - vv and vh of Sentinel-1

train_ds[0]["image"].shape

#Download Clay-1 model

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
ckpt = "https://clay-model-ckpt.s3.amazonaws.com/v0.5.7/mae_v0.5.7_epoch-13_val-loss-0.3098.ckpt"
torch.set_default_device(device)

model = ClayMAEModule.load_from_checkpoint(
    ckpt, metadata_path="model/configs/metadata.yaml", shuffle=False, mask_ratio=0
)
model.eval()

model = model.to(device)


class EarthQuakeDataset:

    def __init__(self,ds):
        self.ds = ds
    
    def __len__(self):
        return len(self.ds)

    def __getitem__(self,idx):
        pre_image = self.ds[idx]["image"][:2,:,:]
        post_image = self.ds[idx]["image"][2:,:,:]
        label = self.ds[idx]["label"]

        sample = {
            "pixels1": pre_image,
            "pixels2": post_image,
            "time": torch.zeros(4),  # Placeholder for time information
            "latlon": torch.zeros(4),  # Placeholder for latlon information
            "label":label
        }

       
        return sample

train_dataset = EarthQuakeDataset(train_ds)

#Generate embeddings

from tqdm.auto import tqdm

gsd =  torch.tensor(10, device=device)
waves = torch.tensor([3.5,4.0], device=device)

embeddings1 = []
embeddings2 = []
target = []

for bid,batch in enumerate(tqdm(train_dl)):
   
    datacube1 = {
        "pixels": batch["pixels1"].to(device),
        "time":batch["time"].to(device),
        "latlon":batch["latlon"].to(device),
        "gsd": gsd,
        "waves": waves,
    }

    datacube2 = {
        "pixels": batch["pixels2"].to(device),
        "time":batch["time"].to(device),
        "latlon":batch["latlon"].to(device),
        "gsd": gsd,
        "waves": waves,
    }

    with torch.no_grad():
        unmsk_patch1, unmsk_idx1, msk_idx1, msk_matrix1 = model.model.encoder(datacube1)
        unmsk_patch2, unmsk_idx2, msk_idx2, msk_matrix2 = model.model.encoder(datacube2)

    emb1 = unmsk_patch1[:, 0, :].cpu().numpy()
    emb2 = unmsk_patch2[:, 0, :].cpu().numpy()

    embeddings1.append(emb1)
    embeddings2.append(emb2)
    target.append(batch["label"].cpu().numpy())

np.save("train_emb.npy",np.concatenate((np.concatenate(embeddings1),np.concatenate(embeddings2)),axis=1))
np.save("train_label.npy",np.concatenate(target))

X = np.load("train_emb.npy")
y = np.load("train_label.npy")

X.shape,y.shape

#((2266, 1536), (2266,))

#Train a RandomForestClassifier on the embeddings

X_train, X_test, y_train, y_test = train_test_split(X, y,random_state=42)

clf = RandomForestClassifier()
clf.fit(X_train, y_train)
predictions = clf.predict(X_test)
cm = confusion_matrix(y_test, predictions, labels=clf.classes_)
disp = ConfusionMatrixDisplay(confusion_matrix=cm,display_labels=clf.classes_)
disp.plot()
plt.show()

![image](https://github.com/user-attachments/assets/55d0dc7f-bf30-461f-9229-fa85150cf16b)

#check F1-score

#Fore reference the competition "ECML-PKDD 2024 - SMAC: Seismic Monitoring and Analysis Challenge" https://www.codabench.org/competitions/2222/#/pages-tab had F1-score upward of 0.908 with training Deep Learning Models

from sklearn.metrics import f1_score

score = f1_score(y_test,predictions)

print("F1 score: ",score)

#F1 score:  0.9088050314465408

[srmsoumya]

Thanks, @anindabitm! It's great to hear this. We've done minimal testing with Sentinel-1 data, so your demo is incredibly helpful.

Would you be open to creating a tutorial for this on Clay docs? I'd be happy to support in any way I can.

[anindabitm]

Thanks @srmsoumya for your response. It will be great opportunity for me and I look forward to create the tutorial under your guidance. I will put together a notebook and please let me know how to connect with you for further guidance on the same

[srmsoumya]

Perfect, once you have the notebook, please create a PR on the clay model repo & we can take it from there.

[anindabitm]

Hi, Many thanks for your reply. I am sorry that I made a mistake while making the PR. Hopefully, that's corrected now. Can you kindly review.