| """Run multi-temporal reconstruction and embedding inference.""" |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.prithvi_eo import PrithviEO2 |
| from train import PrithviDataset, device_from_config |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| device = device_from_config(config) |
| checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=True) |
| if checkpoint["format_version"] != config["data"]["format_version"]: |
| raise ValueError("checkpoint and data formats are incompatible") |
| model = PrithviEO2(checkpoint["model_config"]).to(device) |
| model.load_state_dict(checkpoint["model"]) |
| model.eval() |
| dataset = PrithviDataset(ROOT / config["data"]["root"] / "test.npz", config) |
| pixels = torch.stack([dataset[index]["pixels"] for index in range(len(dataset))]).to(device) |
| temporal = torch.stack([dataset[index]["temporal"] for index in range(len(dataset))]).to(device) |
| location = torch.stack([dataset[index]["location"] for index in range(len(dataset))]).to(device) |
| torch.manual_seed(int(config["seed"])) |
| with torch.no_grad(): |
| output = model(pixels, temporal, location) |
| cls_embedding, patch_embeddings = model.encode(pixels, temporal, location) |
| target = ROOT / config["paths"]["inference_dir"] / "predictions.npz" |
| target.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed( |
| target, |
| format_version=np.asarray(config["data"]["format_version"]), |
| pixels=pixels.cpu().numpy(), |
| reconstruction=output["reconstruction"].cpu().numpy(), |
| mask=output["mask"].cpu().numpy(), |
| embedding=cls_embedding.cpu().numpy(), |
| patch_embeddings=patch_embeddings.cpu().numpy(), |
| temporal_coords=temporal.cpu().numpy(), |
| location_coords=location.cpu().numpy(), |
| class_target=dataset.data["class_target"], |
| regression_target=dataset.data["regression_target"], |
| masked_patch_mse=np.asarray(float(output["loss"])), |
| ) |
| print(f"predictions={target.relative_to(ROOT)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
