Inference on a single Sentinel - 2 patch and visualise the embedding.

[eosurfer]

Hi all,
I have been trying to do inference on a single patch of sentinel - 2.

Bands = ["B02","B03", "B04","B05","B06","B07", "B08", "B8A", "B11", "B12"]
timestamp = 1

Converting to galileo's input

s2 = torch.tensor(s2.to_numpy().transpose(1,2,0)).unsqueeze(2).float() # adding time dimension and converting to float
mnxy = 100
s2 = s2[:mnxy, :mnxy, :] # taking a 100x100px patch
masked_output = construct_galileo_input(s2=s2, normalize=True) 

Inference

with torch.no_grad():
  _embd = model(
            masked_output.space_time_x.unsqueeze(dim=0).float(),
            masked_output.space_x.unsqueeze(dim=0).float(),
            masked_output.time_x.unsqueeze(dim=0).float(),
            masked_output.static_x.unsqueeze(dim=0).float(),
            masked_output.space_time_mask.unsqueeze(dim=0),
            masked_output.space_mask.unsqueeze(dim=0),
            torch.ones_like(masked_output.time_mask.unsqueeze(dim=0)),
            torch.ones_like(masked_output.static_mask.unsqueeze(dim=0)),
            torch.tensor([10]).unsqueeze(dim=0),
            patch_size=2
            )[:-1]

Now I want to visualize these embedding, for that I had look at .averge_token function and I made a custom collapsing function. Because I wanted to preserve the (height, width)

def collapse_and_combine_tc(
        s_t_x: torch.Tensor,
        sp_x: torch.Tensor,
        t_x: torch.Tensor,
        st_x: torch.Tensor,
        s_t_m: torch.Tensor,
        sp_m: torch.Tensor,
        t_m: torch.Tensor,
        st_m: torch.Tensor,
    ):
        s_t_x = rearrange(s_t_x, "b h w t c_g d -> b h w (t c_g) d")
        s_t_m = rearrange(s_t_m, "b h w t c_g-> b h w (t c_g)")
        x = torch.cat(
            [
                s_t_x,
                sp_x,
            ],
            dim=3,
        )
        m = torch.cat([s_t_m, sp_m], dim=3)
        return x, m

collapsed_x, collapsed_mask = collapse_and_combine_tc(*embd)
embd_collapsed = collapsed_x.mean(dim=3)[0] 

Visualize

But when I visualize these embedding after PCA I get this.

from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
embeddings_pca = PCA(n_components=3).fit_transform(embd_collapsed.view(-1,128))
plt.imshow(embeddings_pca.reshape(50,50,3))

Can you help me with this and let me know what am I doing wrong?

And the code in visualizing_embeddings.ipynb produces embedding pixels wise. will it have the spatial information?