Upload 5 files

.gitattributes

@@ -36,3 +36,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 NonReference_RGBN_x4/example_data.safetensor filter=lfs diff=lfs merge=lfs -text
 NonReference_RGBN_x4/hard_constraint.safetensor filter=lfs diff=lfs merge=lfs -text
 NonReference_RGBN_x4/model.safetensor filter=lfs diff=lfs merge=lfs -text
+Reference_RSWIR_x2/example_data.safetensor filter=lfs diff=lfs merge=lfs -text
+Reference_RSWIR_x2/model.safetensor filter=lfs diff=lfs merge=lfs -text

Reference_RSWIR_x2/example_data.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34d594d36a1220c526a55ce086af5f3f6edfd73f4dd966464bfc8b08220854ac
+size 655440

Reference_RSWIR_x2/load.py

@@ -0,0 +1,127 @@
+import pathlib
+import safetensors.torch
+import matplotlib.pyplot as plt
+from sen2sr.models.opensr_baseline.swin import Swin2SR
+from sen2sr.models.tricks import HardConstraint
+from sen2sr.referencex2 import srmodel
+
+# MLSTAC API -----------------------------------------------------------------------
+def example_data(path: pathlib.Path, *args, **kwargs):
+    data_f = path / "example_data.safetensor"    
+    sample = safetensors.torch.load_file(data_f)
+    return  sample["lr"]
+
+def trainable_model(path, device: str = "cpu", *args, **kwargs):
+    trainable_f = path / "model.safetensor"
+
+    # Load model parameters
+    sr_model_weights = safetensors.torch.load_file(trainable_f)
+    params = {
+        "img_size": (64, 64),
+        "in_channels": 10,
+        "out_channels": 6,
+        "embed_dim": 192,
+        "depths": [8] * 8,
+        "num_heads": [8] * 8,
+        "window_size": 4,
+        "mlp_ratio": 4.0,
+        "upscale": 1,
+        "resi_connection": "1conv",
+        "upsampler": "pixelshuffle",
+    }
+    sr_model = Swin2SR(**params)
+    sr_model.load_state_dict(sr_model_weights)
+    sr_model.to(device)
+
+    # Load HardConstraint
+    hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor")
+    hard_constraint = HardConstraint(
+        low_pass_mask=hard_constraint_weights["weights"].to(device),
+        bands= [0, 1, 2, 3, 4, 5],
+        device=device
+    )
+
+    return srmodel(sr_model=sr_model, hard_constraint=hard_constraint, device=device)
+
+
+def compiled_model(path, device: str = "cpu", *args, **kwargs):
+    trainable_f = path / "model.safetensor"
+
+    # Load model parameters
+    sr_model_weights = safetensors.torch.load_file(trainable_f)
+    params = {
+        "img_size": (64, 64),
+        "in_channels": 10,
+        "out_channels": 6,
+        "embed_dim": 192,
+        "depths": [8] * 8,
+        "num_heads": [8] * 8,
+        "window_size": 4,
+        "mlp_ratio": 4.0,
+        "upscale": 1,
+        "resi_connection": "1conv",
+        "upsampler": "pixelshuffle",
+    }
+    sr_model = Swin2SR(**params)
+    sr_model.load_state_dict(sr_model_weights)
+    sr_model = sr_model.to(device)
+    sr_model = sr_model.eval()
+    for param in sr_model.parameters():
+        param.requires_grad = False 
+
+    # Load HardConstraint
+    hard_constraint_weights = safetensors.torch.load_file(path / "hard_constraint.safetensor")
+    hard_constraint = HardConstraint(
+        low_pass_mask=hard_constraint_weights["weights"].to(device),
+        bands= [0, 1, 2, 3, 4, 5],
+        device=device
+    )
+    hard_constraint = hard_constraint.eval()
+    for param in hard_constraint.parameters():
+        param.requires_grad = False
+
+    return srmodel(sr_model=sr_model, hard_constraint=hard_constraint, device=device)
+
+
+def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs):
+    # Load model
+    model = compiled_model(path, device)
+
+    # Load data
+    lr = example_data(path)
+
+    # Run model
+    sr = model(lr.to(device))
+
+    # Create the viz
+    lr_rgb = lr[0, [2, 1, 0]].cpu().numpy().transpose(1, 2, 0)
+    sr_rgb = sr[0, [2, 1, 0]].cpu().numpy().transpose(1, 2, 0)
+    
+    lr_swirs = lr[0, [9, 8, 7]].cpu().numpy().transpose(1, 2, 0)
+    sr_swirs = sr[0, [9, 8, 7]].cpu().numpy().transpose(1, 2, 0)
+       
+    lr_reds = lr[0, [6, 5, 4]].cpu().numpy().transpose(1, 2, 0)
+    sr_reds = sr[0, [6, 5, 4]].cpu().numpy().transpose(1, 2, 0)
+    
+
+    #Display results
+    lr_slice = slice(16, 32+80)
+    hr_slice = slice(lr_slice.start*1, lr_slice.stop*1)
+    fig, ax = plt.subplots(3, 2, figsize=(8, 12))
+    ax = ax.flatten()
+    ax[0].imshow(lr_rgb[lr_slice]*2)
+    ax[0].set_title("LR RGB")
+    ax[1].imshow(sr_rgb[hr_slice]*2)
+    ax[1].set_title("SR RGB")
+    ax[2].imshow(lr_swirs[lr_slice]*2)
+    ax[2].set_title("LR SWIR")
+    ax[3].imshow(sr_swirs[hr_slice]*2)
+    ax[3].set_title("SR SWIR")
+    ax[4].imshow(lr_reds[lr_slice]*2)
+    ax[4].set_title("LR RED")
+    ax[5].imshow(sr_reds[hr_slice]*2)
+    ax[5].set_title("SR RED")
+    for a in ax:
+        a.axis("off")
+    fig.tight_layout()
+    return fig

Reference_RSWIR_x2/mlm.json

@@ -0,0 +1,187 @@
+{
+    "type": "Feature",
+    "stac_version": "1.1.0",
+    "stac_extensions": [
+        "https://stac-extensions.github.io/mlm/v1.4.0/schema.json"
+    ],
+    "id": "SPAN_model",
+    "geometry": {
+        "type": "Polygon",
+        "coordinates": [
+            [
+                [
+                    -180.0,
+                    -90.0
+                ],
+                [
+                    -180.0,
+                    90.0
+                ],
+                [
+                    180.0,
+                    90.0
+                ],
+                [
+                    180.0,
+                    -90.0
+                ],
+                [
+                    -180.0,
+                    -90.0
+                ]
+            ]
+        ]
+    },
+    "bbox": [
+        -180,
+        -90,
+        180,
+        90
+    ],
+    "properties": {
+        "start_datetime": "1900-01-01T00:00:00Z",
+        "end_datetime": "9999-01-01T00:00:00Z",
+        "description": "A Swift Parameter-free Attention Network (SPAN) trained on the SEN2NAIPv2 dataset to enhance RSWIRs Sentinel-2 bands, improving spatial resolution from 20 meters to 10 meters.",
+        "forward_backward_pass": {
+            "32": 24.957696,
+            "64": 97.858304,
+            "128": 387.576576,
+            "256": 1542.681344,
+            "512": 6155.563776
+        },
+        "dependencies": [
+            "torch",
+            "safetensors.torch",
+            "sen2sr"
+        ],
+        "mlm:framework": "pytorch",
+        "mlm:framework_version": "2.1.2+cu121",
+        "file:size": 1889984,
+        "mlm:memory_size": 1,
+        "mlm:accelerator": "cuda",
+        "mlm:accelerator_constrained": false,
+        "mlm:accelerator_summary": "Unknown",
+        "mlm:name": "CNN_Light_SR",
+        "mlm:architecture": "SPAN",
+        "mlm:tasks": [
+            "super-resolution"
+        ],
+        "mlm:input": [
+            {
+                "name": "10 Band Sentinel-2",
+                "bands": [
+                    "B04",
+                    "B03",
+                    "B02",
+                    "B08",
+                    "B05",
+                    "B06",
+                    "B07",
+                    "B8A",
+                    "B11",
+                    "B12"
+                ],
+                "input": {
+                    "shape": [
+                        -1,
+                        10,
+                        128,
+                        128
+                    ],
+                    "dim_order": [
+                        "batch",
+                        "channel",
+                        "height",
+                        "width"
+                    ],
+                    "data_type": "float16"
+                },
+                "pre_processing_function": null
+            }
+        ],
+        "mlm:output": [
+            {
+                "name": "super-resolution",
+                "bands": [
+                    "B05",
+                    "B06",
+                    "B07",
+                    "B8A",
+                    "B11",
+                    "B12"
+                ],
+                "tasks": [
+                    "super-resolution"
+                ],
+                "result": {
+                    "shape": [
+                        -1,
+                        6,
+                        128,
+                        128
+                    ],
+                    "dim_order": [
+                        "batch",
+                        "channel",
+                        "height",
+                        "width"
+                    ],
+                    "data_type": "float16"
+                },
+                "classification:classes": [],
+                "post_processing_function": null
+            }
+        ],
+        "mlm:total_parameters": 472496,
+        "mlm:pretrained": true,
+        "datetime": null
+    },
+    "links": [],
+    "assets": {
+        "trainable": {
+            "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SR/Reference_RSWIR_x2/model.safetensor",
+            "type": "application/octet-stream; application=safetensor",
+            "title": "Pytorch model weights checkpoint",
+            "description": "The weights of the model in safetensor format.",
+            "mlm:artifact_type": "safetensor.torch.save_file",
+            "roles": [
+                "mlm:model",
+                "mlm:weights",
+                "data"
+            ]
+        },
+        "hardconstraint": {
+            "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SR/Reference_RSWIR_x2/hard_constraint.safetensor",
+            "type": "application/octet-stream; application=safetensor",
+            "title": "Pytorch hard constraint weights checkpoint. It is used to load faster the hard constraint module.",
+            "description": "The weights of the model in safetensor format.",
+            "mlm:artifact_type": "safetensor.torch.save_file",
+            "roles": [
+                "mlm:model",
+                "mlm:weights",
+                "data"
+            ]
+        },
+        "source_code": {
+            "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SR/Reference_RSWIR_x2/load.py",
+            "type": "text/x-python",
+            "title": "Model load script",
+            "description": "Python script to load the model.",
+            "roles": [
+                "mlm:source_code",
+                "code"
+            ]
+        },
+        "example_data": {
+            "href": "https://huggingface.co/tacofoundation/SEN2SR/resolve/main/SEN2SR/Reference_RSWIR_x2/example_data.safetensor",
+            "type": "application/octet-stream; application=safetensors",
+            "title": "Example Sentinel-2 image",
+            "description": "Example Sentinel-2 image for model inference.",
+            "roles": [
+                "mlm:example_data",
+                "data"
+            ]
+        }
+    },
+    "collection": "ml-model"
+}

Reference_RSWIR_x2/model.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d1b93248ca0bb24cc20a6c2e83b2acd70517357fa322f9051af8f378d54eecc
+size 137790600