import torch
from einops import rearrange
from torchgeo.models import RCF
from sklearn.decomposition import PCA
# Each sample in a batch is a time-series of images
# batch = torch.randn(8, 4, 27, 64, 64)  # (b, t, c, h, w)
dataset = ...  # a torch dataset that returns dict(image=...)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False)
# Initialize the model
model = RCF(in_channels=27, features=4096, kernel_size=3, mode="gaussian")
model.eval()
# Loop over your images and collect all embeddings
embeddings = []
for batch in dataloader:
    b, t = batch.shape[0], batch.shape[1]
    with torch.inference_mode():
        # Run MOSAIKS across all imagery indepedently
        x = rearrange(batch, "b t c h w -> (b t) c h w")
        emb = model(x)
        
        # Average pool over the time dimension
        emb = rearrange(emb, "(b t) c -> b t c", b=b, t=t)
        emb = emb.mean(dim=1)
        embeddings.append(emb)
# Perform PCA on the embeddings
embeddings = torch.cat(embeddings, dim=0).numpy()
embeddings = PCA(n_components=1024).fit_transform(embeddings)