Update README.md

README.md

@@ -14,5 +14,76 @@ max_batch_size: 64
 merge_mode: weighted_average
 ---
 
-# Model Card
+First run the following to setup the environment and get the official model code
+
+```bash
+# Clone the official repo
+git clone git@github.com:facebookresearch/HighResCanopyHeight.git
+
+# Install dependencies
+pip install stac-model[torch]
+
+# Download the official pretrained checkpoints
+mkdir checkpoints && aws s3 --no-sign-request sync s3://dataforgood-fb-data/forests/v1/models/saved_checkpoints/ checkpoints/
+```
+
+Export the model using the following:
+
+```python
+from pathlib import Path
+import sys
+sys.path.append("HighResCanopyHeight")
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.v2 as T
+from stac_model.torch.export import export, package
+
+import src.transforms
+from inference import SSLAE
+
+
+# Create model and load checkpoint
+class TreeCanopyHeightModel(nn.Module):
+    def __init__(self, classify=True, huge=True):
+        super().__init__()
+        self.model = SSLAE(pretrained=None, classify=classify, huge=huge, n_bins=256)
+
+    def forward(self, x):
+        outputs = self.model(x)
+        pred = 10 * outputs + 0.001
+        return pred.relu()
+
+path = "checkpoints/SSLhuge_satellite.pth"
+ckpt = torch.load(path, map_location="cpu", weights_only=False)
+state_dict = {f"model.{k}": v for k, v in ckpt["state_dict"].items()}
+model = TreeCanopyHeightModel()
+model.load_state_dict(state_dict)
+
+# Create exportable transforms
+original_transform = src.transforms.SSLNorm().Trans
+norm = original_transform.transforms[-1]
+
+transforms = nn.Sequential(
+    T.Normalize(mean=[0], std=[255]),  # replace ToTensor() with normalize to 0-1
+    T.Normalize(mean=norm.mean, std=norm.std)
+)
+
+# Export and save to pt2
+model_program, transforms_program = export(
+    input_shape=[-1, 3, 224, 224],
+    model=model,
+    transforms=transforms,
+    device="cpu",
+    dtype=torch.float32,
+)
+package(
+    output_file=Path("model.pt2"),
+    model_program=model_program,
+    transforms_program=transforms_program,
+    metadata_properties=None,
+    aoti_compile_and_package=False
+)
+```
+