Prithvi 1.0 still supported?
Hello! My team and I are trying to run Prithvi 1.0 to do reconstruction. We've been trying to use reconstruction to handle images with lots of clouds in them.
This notebook goes through what we're trying to do: https://github.com/NASA-IMPACT/hls-foundation-os/blob/main/exploration.ipynb
However it seems like perhaps support has shifted, especially with the introduction of Terratorch. Is Prithvi 1.0 on Terratorch? We're specifically facing an issue on this line:
It doesn't seem like there's any examples of Prithvi 2.0 doing reconstruction. But the example there is of Prithvi 1.0 doesn't run anymore and has an error at that line.
If anyone has any insights, please help!
@tishyac3141 A lot changed in the latest TerraTorch version and it is not much simpler to build the MAE model. You can initialize it now with
from terratorch.registry import FULL_MODEL_REGISTRY
model = FULL_MODEL_REGISTRY.build('prithvi_eo_v1_100_mae', pretrained=True)
A small pipeline could look like this:
import torch
from terratorch.registry import FULL_MODEL_REGISTRY
from terratorch.models.backbones.prithvi_vit import PRITHVI_V1_MEAN, PRITHVI_V1_STD
# Load MAE model via TerraTorch
model = FULL_MODEL_REGISTRY.build('prithvi_eo_v1_100_mae', pretrained=True)
mean = torch.tensor(PRITHVI_V1_MEAN)
std = torch.tensor(PRITHVI_V1_STD)
# Load your data (using random values as example)
data = torch.rand(1, 6, 224, 224)
data = (data - mean[None, :, None, None]) / std[None, :, None, None]
loss, reconstructed, mask = model(data)
reconstructed = reconstructed * std[None, :, None, None] + mean[None, :, None, None]
However, the code currently does not support providing a custom mask for the reconstruction, so you need to update the forward function (see https://github.com/IBM/terratorch/blob/86224762d4b2128f2673a5eb5bc95d2b2f911063/terratorch/models/backbones/prithvi_mae.py#L812)