Upload 5 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example_data.safetensor filter=lfs diff=lfs merge=lfs -text
+model.safetensor filter=lfs diff=lfs merge=lfs -text

example_data.safetensor

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

load.py

@@ -0,0 +1,89 @@
+import importlib.util
+import pathlib
+
+import matplotlib.pyplot as plt
+import safetensors.torch
+import torch
+
+
+def load_model_module(model_path: pathlib.Path):
+    model_path = model_path.resolve()
+
+    spec = importlib.util.spec_from_file_location("model", model_path)
+    model = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(model)
+
+    return model
+
+
+class GeneratorNormal(torch.nn.Module):
+    def __init__(self, model):
+        super(GeneratorNormal, self).__init__()
+        self.model = model
+
+    def forward(self, X):
+        X = torch.clamp(X, 0, 1)
+        X = torch.nan_to_num(X, nan=1.0)
+        X = X.permute(0, 2, 3, 4, 1).contiguous()
+        return self.model(X)[0].permute(0, 4, 1, 2, 3)
+
+
+def example_data(path: pathlib.Path, *args, **kwargs):
+    data_f = path / "example_data.safetensor"
+    return safetensors.torch.load_file(data_f)["example_data"][None].float()
+
+
+def trainable_model(path, device="cpu", *args, **kwargs):
+    weights = safetensors.torch.load_file(path / "model.safetensor")
+    model = load_model_module(path / "model.py").Generator(
+        device=device, inputChannels=4, outputChannels=4
+    )
+    model.load_state_dict(weights)
+    return model
+
+
+def compiled_model(path, device="cpu", *args, **kwargs):
+    weights = safetensors.torch.load_file(path / "model.safetensor")
+    model = load_model_module(path / "model.py").Generator(
+        device=device, inputChannels=4, outputChannels=4
+    )
+    model.load_state_dict(weights)
+    model = model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+    return GeneratorNormal(model.to(device))
+
+
+def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs):
+    # Load model
+    model = compiled_model(path, device)
+
+    # Load data
+    s2_ts = example_data(path)
+
+    # Run model
+    gap_filled = model(s2_ts.to(device))
+
+    # Convert to CPU and detach for plotting
+    s2_ts = s2_ts.squeeze(0).detach().cpu()  # [T, C, H, W]
+    gap_filled = gap_filled.squeeze(0).detach().cpu()
+
+    num_timesteps = s2_ts.shape[0]
+    rgb_indices = [2, 1, 0]  # Assuming RGB is BGR in channel order (4 bands)
+
+    fig, axs = plt.subplots(2, num_timesteps, figsize=(3 * num_timesteps, 6))
+
+    for t in range(num_timesteps):
+        original_rgb = s2_ts[t, rgb_indices].permute(1, 2, 0).clamp(0, 1).numpy()
+        filled_rgb = gap_filled[t, rgb_indices].permute(1, 2, 0).clamp(0, 1).numpy()
+
+        axs[0, t].imshow(original_rgb * 3)
+        axs[0, t].axis("off")
+        axs[0, t].set_title(f"Original t={t}")
+
+        axs[1, t].imshow(filled_rgb * 3)
+        axs[1, t].axis("off")
+        axs[1, t].set_title(f"Filled t={t}")
+
+    plt.tight_layout()
+    return fig

mlm.json

@@ -0,0 +1,203 @@
+{
+    "type": "Feature",
+    "stac_version": "1.1.0",
+    "stac_extensions": [
+        "https://stac-extensions.github.io/mlm/v1.4.0/schema.json"
+    ],
+    "id": "UNetMobV2_V2 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 UNet model trained on Sentinel-2 imagery for cloud segmentation.",
+        "forward_backward_pass": {
+            "32": 3.229184,
+            "64": 12.916736,
+            "128": 51.666944,
+            "256": 206.667776,
+            "512": 826.671104,
+            "1024": 3306.684416,
+            "2048": 13226.737664
+        },
+        "dependencies": [
+            "torch",
+            "segmentation-models-pytorch",
+            "safetensors.torch"
+        ],
+        "mlm:framework": "pytorch",
+        "mlm:framework_version": "2.1.2+cu121",
+        "file:size": 26529040,
+        "mlm:memory_size": 1,
+        "mlm:accelerator": "cuda",
+        "mlm:accelerator_constrained": false,
+        "mlm:accelerator_summary": "Unknown",
+        "mlm:name": "UNetMobV2_V1",
+        "mlm:architecture": "UNetMobV2",
+        "mlm:tasks": [
+            "semantic-segmentation"
+        ],
+        "mlm:input": [
+            {
+                "name": "13 Band Sentinel-2 Batch",
+                "bands": [
+                    "B01",
+                    "B02",
+                    "B03",
+                    "B04",
+                    "B05",
+                    "B06",
+                    "B07",
+                    "B08",
+                    "B8A",
+                    "B09",
+                    "B10",
+                    "B11",
+                    "B12"
+                ],
+                "input": {
+                    "shape": [
+                        -1,
+                        13,
+                        512,
+                        512
+                    ],
+                    "dim_order": [
+                        "batch",
+                        "channel",
+                        "height",
+                        "width"
+                    ],
+                    "data_type": "float32"
+                },
+                "pre_processing_function": null
+            }
+        ],
+        "mlm:output": [
+            {
+                "name": "semantic-segmentation",
+                "tasks": [
+                    "semantic-segmentation"
+                ],
+                "result": {
+                    "shape": [
+                        -1,
+                        4,
+                        512,
+                        512
+                    ],
+                    "dim_order": [
+                        "batch",
+                        "channel",
+                        "height",
+                        "width"
+                    ],
+                    "data_type": "float32"
+                },
+                "classification:classes": [
+                    {
+                        "value": 0,
+                        "name": "Clear",
+                        "description": "Clear"
+                    },
+                    {
+                        "value": 1,
+                        "name": "Thick Clouds",
+                        "description": "Thick Clouds"
+                    },
+                    {
+                        "value": 2,
+                        "name": "Thin Clouds",
+                        "description": "Thin Clouds"
+                    },
+                    {
+                        "value": 3,
+                        "name": "Cloud Shadows",
+                        "description": "Cloud Shadows"
+                    }
+                ],
+                "post_processing_function": null
+            }
+        ],
+        "mlm:total_parameters": 6632260,
+        "mlm:pretrained": true,
+        "datetime": null
+    },
+    "links": [],
+    "assets": {
+        "trainable": {
+            "href": "https://huggingface.co/tacofoundation/GANFilling/resolve/main/model.safetensor",
+            "type": "application/octet-stream; application=safetensor",
+            "title": "Pytorch weights checkpoint",
+            "description": "A UNet model trained on Sentinel-2 imagery for cloud segmentation.The model was trained using the CloudSEN12 dataset.",
+            "mlm:artifact_type": "safetensor.torch.save_file",
+            "roles": [
+                "mlm:model",
+                "mlm:weights",
+                "data"
+            ]
+        },
+        "source_code": {
+            "href": "https://huggingface.co/tacofoundation/GANFilling/resolve/main/load.py",
+            "type": "text/x-python",
+            "title": "Model load script",
+            "description": "Source code to run the model.",
+            "roles": [
+                "mlm:source_code",
+                "code"
+            ]
+        },
+        "source_code_model": {
+            "href": "https://huggingface.co/tacofoundation/GANFilling/resolve/main/model.py",
+            "type": "text/x-python",
+            "title": "Model load script",
+            "description": "Source code to run the model.",
+            "roles": [
+                "mlm:source_code",
+                "code"
+            ]
+        },        
+        "example_data": {
+            "href": "https://huggingface.co/tacofoundation/GANFilling/resolve/main/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": "GANFilling"
+}

model.py

@@ -0,0 +1,348 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ConvLSTMCell(nn.Module):
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias, device):
+        """
+        Initialize ConvLSTM cell.
+
+        Parameters
+        ----------
+        input_dim: int
+            Number of channels of input tensor.
+        hidden_dim: int
+            Number of channels of hidden state.
+        kernel_size: (int, int)
+            Size of the convolutional kernel.
+        bias: bool
+            Whether or not to add the bias.
+        """
+
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.kernel_size = kernel_size
+        self.padding = kernel_size[0] // 2, kernel_size[1] // 2
+        self.bias = bias
+        self.device = device
+
+        self.conv = nn.Conv2d(
+            in_channels=self.input_dim + self.hidden_dim,
+            out_channels=4 * self.hidden_dim,
+            kernel_size=self.kernel_size,
+            padding=self.padding,
+            bias=self.bias,
+        )
+
+    def __initStates(self, size):
+        return torch.zeros(size).to(self.device), torch.zeros(size).to(self.device)
+        # return torch.zeros(size).cuda(), torch.zeros(size).cuda()
+
+    def forward(self, input_tensor, cur_state):
+        if cur_state == None:
+            h_cur, c_cur = self.__initStates(
+                [
+                    input_tensor.shape[0],
+                    self.hidden_dim,
+                    input_tensor.shape[2],
+                    input_tensor.shape[3],
+                ]
+            )
+        else:
+            h_cur, c_cur = cur_state
+
+        combined = torch.cat(
+            [input_tensor, h_cur], dim=1
+        )  # concatenate along channel axis
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+
+        return h_next, c_next
+
+    def init_hidden(self, batch_size, image_size):
+        height, width = image_size
+        return (
+            torch.zeros(
+                batch_size,
+                self.hidden_dim,
+                height,
+                width,
+                device=self.conv.weight.device,
+            ),
+            torch.zeros(
+                batch_size,
+                self.hidden_dim,
+                height,
+                width,
+                device=self.conv.weight.device,
+            ),
+        )
+
+
+class ConvLSTM(nn.Module):
+    """
+
+    Parameters:
+        input_dim: Number of channels in input
+        hidden_dim: Number of hidden channels
+        kernel_size: Size of kernel in convolutions
+        num_layers: Number of LSTM layers stacked on each other
+        batch_first: Whether or not dimension 0 is the batch or not
+        bias: Bias or no bias in Convolution
+        return_all_layers: Return the list of computations for all layers
+        Note: Will do same padding.
+
+    Input:
+        A tensor of size B, T, C, H, W or T, B, C, H, W
+    Output:
+        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).
+            0 - layer_output_list is the list of lists of length T of each output
+            1 - last_state_list is the list of last states
+                    each element of the list is a tuple (h, c) for hidden state and memory
+    Example:
+        >> x = torch.rand((32, 10, 64, 128, 128))
+        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)
+        >> _, last_states = convlstm(x)
+        >> h = last_states[0][0]  # 0 for layer index, 0 for h index
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        kernel_size,
+        num_layers,
+        batch_first=False,
+        bias=True,
+        return_all_layers=False,
+    ):
+        super(ConvLSTM, self).__init__()
+
+        self._check_kernel_size_consistency(kernel_size)
+
+        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
+        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
+        if not len(kernel_size) == len(hidden_dim) == num_layers:
+            raise ValueError("Inconsistent list length.")
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+
+        cell_list = []
+        for i in range(0, self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+
+            cell_list.append(
+                ConvLSTMCell(
+                    input_dim=cur_input_dim,
+                    hidden_dim=self.hidden_dim[i],
+                    kernel_size=self.kernel_size[i],
+                    bias=self.bias,
+                )
+            )
+
+        self.cell_list = nn.ModuleList(cell_list)
+
+    def forward(self, input_tensor, hidden_state=None):
+        """
+
+        Parameters
+        ----------
+        input_tensor: todo
+            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)
+        hidden_state: todo
+            None. todo implement stateful
+
+        Returns
+        -------
+        last_state_list, layer_output
+        """
+        if not self.batch_first:
+            # (t, b, c, h, w) -> (b, t, c, h, w)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+
+        b, _, _, h, w = input_tensor.size()
+
+        # Implement stateful ConvLSTM
+        if hidden_state is not None:
+            raise NotImplementedError()
+        else:
+            # Since the init is done in forward. Can send image size here
+            hidden_state = self._init_hidden(batch_size=b, image_size=(h, w))
+
+        layer_output_list = []
+        last_state_list = []
+
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+
+        for layer_idx in range(self.num_layers):
+
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            for t in range(seq_len):
+                h, c = self.cell_list[layer_idx](
+                    input_tensor=cur_layer_input[:, t, :, :, :], cur_state=[h, c]
+                )
+                output_inner.append(h)
+
+            layer_output = torch.stack(output_inner, dim=1)
+            cur_layer_input = layer_output
+
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+
+        return layer_output_list, last_state_list
+
+    def _init_hidden(self, batch_size, image_size):
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))
+        return init_states
+
+    @staticmethod
+    def _check_kernel_size_consistency(kernel_size):
+        if not (
+            isinstance(kernel_size, tuple)
+            or (
+                isinstance(kernel_size, list)
+                and all([isinstance(elem, tuple) for elem in kernel_size])
+            )
+        ):
+            raise ValueError("`kernel_size` must be tuple or list of tuples")
+
+    @staticmethod
+    def _extend_for_multilayer(param, num_layers):
+        if not isinstance(param, list):
+            param = [param] * num_layers
+        return param
+
+
+def normal_init(m, mean, std):
+    if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
+        m.weight.data.normal_(mean, std)
+        m.bias.data.zero_()
+
+
+class Generator(nn.Module):
+    def __init__(self, device, inputChannels=4, outputChannels=3, d=64):
+        super().__init__()
+        self.d = d
+        self.device = device
+
+        self.conv1 = nn.Conv2d(inputChannels, d, 3, 2, 1)
+        self.conv2 = nn.Conv2d(d, d * 2, 3, 2, 1)
+        self.conv3 = nn.Conv2d(d * 2, d * 4, 3, 2, 1)
+        self.conv4 = nn.Conv2d(d * 4, d * 8, 3, 2, 1)
+        self.conv5 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+        self.conv6 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+        self.conv7 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+
+        self.conv_lstm_d1 = ConvLSTMCell(d * 8, d * 8, (3, 3), False, device)
+        self.conv_lstm_d2 = ConvLSTMCell(d * 8 * 2, d * 8, (3, 3), False, device)
+        self.conv_lstm_d3 = ConvLSTMCell(d * 8 * 2, d * 8, (3, 3), False, device)
+        self.conv_lstm_d4 = ConvLSTMCell(d * 8 * 2, d * 4, (3, 3), False, device)
+        self.conv_lstm_d5 = ConvLSTMCell(d * 4 * 2, d * 2, (3, 3), False, device)
+        self.conv_lstm_d6 = ConvLSTMCell(d * 2 * 2, d, (3, 3), False, device)
+        self.conv_lstm_d7 = ConvLSTMCell(d * 2, d, (3, 3), False, device)
+
+        self.conv_lstm_e1 = ConvLSTMCell(d, d, (3, 3), False, device)
+        self.conv_lstm_e2 = ConvLSTMCell(d * 2, d * 2, (3, 3), False, device)
+        self.conv_lstm_e3 = ConvLSTMCell(d * 4, d * 4, (3, 3), False, device)
+        self.conv_lstm_e4 = ConvLSTMCell(d * 8, d * 8, (3, 3), False, device)
+        self.conv_lstm_e5 = ConvLSTMCell(d * 8, d * 8, (3, 3), False, device)
+        self.conv_lstm_e6 = ConvLSTMCell(d * 8, d * 8, (3, 3), False, device)
+        self.conv_lstm_e7 = ConvLSTMCell(d * 8, d * 8, (3, 3), False, device)
+
+        self.up = nn.Upsample(scale_factor=2)
+        self.conv_out = nn.Conv2d(d, outputChannels, 3, 1, 1)
+
+        self.slope = 0.2
+
+    def weight_init(self, mean, std):
+        for m in self._modules:
+            normal_init(self._modules[m], mean, std)
+
+    def forward_step(self, input, states_encoder, states_decoder):
+
+        e1 = self.conv1(input)
+        states_e1 = self.conv_lstm_e1(e1, states_encoder[0])
+        e2 = self.conv2(F.leaky_relu(states_e1[0], self.slope))
+        states_e2 = self.conv_lstm_e2(e2, states_encoder[1])
+        e3 = self.conv3(F.leaky_relu(states_e2[0], self.slope))
+        states_e3 = self.conv_lstm_e3(e3, states_encoder[2])
+        e4 = self.conv4(F.leaky_relu(states_e3[0], self.slope))
+        states_e4 = self.conv_lstm_e4(e4, states_encoder[3])
+        e5 = self.conv5(F.leaky_relu(states_e4[0], self.slope))
+        states_e5 = self.conv_lstm_e5(e5, states_encoder[4])
+        e6 = self.conv6(F.leaky_relu(states_e5[0], self.slope))
+        states_e6 = self.conv_lstm_e6(e6, states_encoder[5])
+        e7 = self.conv7(F.leaky_relu(states_e6[0], self.slope))
+
+        states1 = self.conv_lstm_d1(F.relu(e7), states_decoder[0])
+        d1 = self.up(states1[0])
+        d1 = torch.cat([d1, e6], 1)
+
+        states2 = self.conv_lstm_d2(F.relu(d1), states_decoder[1])
+        d2 = self.up(states2[0])
+        d2 = torch.cat([d2, e5], 1)
+
+        states3 = self.conv_lstm_d3(F.relu(d2), states_decoder[2])
+        d3 = self.up(states3[0])
+        d3 = torch.cat([d3, e4], 1)
+
+        states4 = self.conv_lstm_d4(F.relu(d3), states_decoder[3])
+        d4 = self.up(states4[0])
+        d4 = torch.cat([d4, e3], 1)
+
+        states5 = self.conv_lstm_d5(F.relu(d4), states_decoder[4])
+        d5 = self.up(states5[0])
+        d5 = torch.cat([d5, e2], 1)
+
+        states6 = self.conv_lstm_d6(F.relu(d5), states_decoder[5])
+        d6 = self.up(states6[0])
+        d6 = torch.cat([d6, e1], 1)
+
+        states7 = self.conv_lstm_d7(F.relu(d6), states_decoder[6])
+        d7 = self.up(states7[0])
+
+        o = torch.clip(torch.tanh(self.conv_out(d7)), min=-0.0, max=1)
+
+        states_e = [states_e1, states_e2, states_e3, states_e4, states_e5, states_e6]
+        states_d = [states1, states2, states3, states4, states5, states6, states7]
+
+        return o, (states_e, states_d)
+
+    def forward(self, tensor):
+        states_encoder = (None, None, None, None, None, None, None)
+        states_decoder = (None, None, None, None, None, None, None)
+        output = torch.empty_like(tensor)
+        for timeStep in range(tensor.shape[4]):
+            output[:, :, :, :, timeStep], states = self.forward_step(
+                tensor[:, :, :, :, timeStep], states_encoder, states_decoder
+            )
+            states_encoder, states_decoder = states[0], states[1]
+        return output, states

model.safetensor

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