deploy changes

Author: github-actions[bot]

.github/workflows/deploy_to_open.yaml

@@ -0,0 +1,61 @@
+name: Deploy to Open Repo
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+    if: |
+      github.event.repository.name == 'internal_asparagus' &&
+      github.event.push.pusher.username != 'github-actions[bot]'
+    name: Push built output to target repo via SSH
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository (no persisted credentials)
+        uses: actions/checkout@v4
+        with:
+          ref: deploy
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Configure git author
+        run: |
+          git config --global user.name "github-actions[bot]"
+          git config --global user.email "41898282+github-actions[bot]@users.noreply.github.com"
+
+      - name: Update deployment branch and remove this workflow so it doesn't get pushed
+        run: |
+          git merge --squash -X theirs origin/main --allow-unrelated-histories
+          git restore --staged README.md
+          git commit -m "deploy changes"
+          
+      - name: Start ssh-agent and add deploy key
+        uses: webfactory/ssh-agent@v0.9.1
+        with:
+          ssh-private-key: ${{ secrets.DEPLOY }}
+
+      - name: Ensure github.com is in known_hosts
+        run: |
+          mkdir -p ~/.ssh
+          ssh-keyscan github.com >> ~/.ssh/known_hosts
+          chmod 644 ~/.ssh/known_hosts
+
+      - name: Set SSH remote and push via SSH
+        run: |
+          echo "Pushing to SSH remote: git@github.com:Sllambias/asparagus.git -> branch main"
+
+          git remote remove origin || true
+          git remote add ssh-origin git@github.com:Sllambias/asparagus.git
+
+          # Sanity checks
+          git remote -v
+          echo "SSH_AUTH_SOCK=$SSH_AUTH_SOCK"
+          ssh -T -o StrictHostKeyChecking=no git@github.com || true
+
+          git push --force -u ssh-origin HEAD:main

asparagus/modules/lightning_modules/base_module.py

@@ -60,7 +60,7 @@ def __init__(
         self.repeat_stem_weights = repeat_stem_weights
         assert 0 < cosine_period_ratio <= 1
 
-        self.save_hyperparameters(ignore=["model"])
+        self.save_hyperparameters(ignore=["model", "train_transforms", "val_transforms", "test_transforms"])
         self.model = model
 
         if weights is not None:

asparagus/modules/networks/resenc_unet.py

@@ -1,4 +1,95 @@
+import logging
+import torch
+from gardening_tools.modules.networks.BaseNet import BaseNet
+from gardening_tools.modules.networks.components.blocks import ResidualBlock
+from gardening_tools.modules.networks.components.encoders import ResidualUNetEncoder
+from gardening_tools.modules.networks.components.heads import ClsRegHead
 from gardening_tools.modules.networks.resunet import ResidualEncoderUNet
+from torch import nn
+from typing import List, Tuple, Type, Union
+
+
+class ResidualEncoderUNetCLSREG(BaseNet):
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        dimensions: str,
+        kernel_size: int,
+        stride: int,
+        features_per_stage: list,
+        n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+        conv_bias: bool = True,
+        encoder_basic_block: Type[ResidualBlock] = ResidualBlock,
+        decoder: nn.Module = ClsRegHead,
+        norm_op_kwargs={"eps": 1e-05, "affine": True},
+        dropout_op=None,
+        dropout_op_kwargs=None,
+        nonlin=torch.nn.LeakyReLU,
+        nonlin_kwargs={"inplace": True},
+    ):
+        super().__init__()
+
+        # Extract dropout rates from kwargs
+        if dropout_op_kwargs is None:
+            dropout_op_kwargs = {}
+
+        encoder_dropout_rate = dropout_op_kwargs.get("encoder_dropout_rate", 0.0)
+        decoder_dropout_rate = dropout_op_kwargs.get("decoder_dropout_rate", 0.0)
+        inplace = dropout_op_kwargs.get("inplace", True)
+
+        if dimensions == "2D":
+            conv_op = nn.Conv2d
+            dropout_op = nn.Dropout2d
+            norm_op = nn.InstanceNorm2d
+            pool_op = nn.MaxPool2d
+            clsreg_pool_op = nn.AdaptiveAvgPool2d
+        elif dimensions == "3D":
+            conv_op = nn.Conv3d
+            dropout_op = nn.Dropout3d
+            norm_op = nn.InstanceNorm3d
+            pool_op = nn.MaxPool3d
+            clsreg_pool_op = nn.AdaptiveAvgPool3d
+        else:
+            logging.warning("Uuh, dimensions not in ['2D', '3D']")
+
+        self.num_classes = output_channels
+
+        self.stem_weight_name = "encoder.stem.conv1.conv.weight"
+
+        self.encoder = ResidualUNetEncoder(
+            input_channels=input_channels,
+            features_per_stage=features_per_stage,
+            conv_op=conv_op,
+            kernel_size=kernel_size,
+            stride=stride,
+            n_blocks_per_stage=n_blocks_per_stage,
+            conv_bias=conv_bias,
+            norm_op=norm_op,
+            norm_op_kwargs=norm_op_kwargs,
+            dropout_op=dropout_op,
+            dropout_op_kwargs={"p": encoder_dropout_rate, "inplace": inplace},
+            nonlin=nonlin,
+            nonlin_kwargs=nonlin_kwargs,
+            block=encoder_basic_block,
+            pool_op=pool_op,
+        )
+
+        self.decoder = decoder(
+            pool_op=clsreg_pool_op,
+            input_channels=features_per_stage[-1],
+            output_channels=output_channels,
+            dropout_rate=decoder_dropout_rate,
+        )
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def forward_with_features(self, x):
+        skips = self.encoder(x)
+        output = self.decoder(skips)
+        return output, skips[-1]
 
 
 # Encoder 29M parameters
@@ -46,6 +137,26 @@ def resenc_unet_b(
     )
 
 
+# Encoder 90M parameters
+# Full model - 90.3 M    Total params
+def resenc_unet_b_clsreg(
+    input_channels: int = 1,
+    output_channels: int = 1,
+    dimensions: str = "3D",
+    dropout_op_kwargs: dict = None,
+):
+    return ResidualEncoderUNetCLSREG(
+        dimensions=dimensions,
+        input_channels=input_channels,
+        output_channels=output_channels,
+        features_per_stage=(32, 64, 128, 256, 320, 320),
+        stride=2,
+        kernel_size=3,
+        n_blocks_per_stage=(1, 3, 4, 6, 6, 6),
+        dropout_op_kwargs=dropout_op_kwargs,
+    )
+
+
 # Encoder 345M parameters
 # Full model 391M parameters
 def resenc_unet_l(

asparagus/pipeline/auto_configuration/versioning.py

@@ -26,6 +26,7 @@ def pathing(cfg, train=True):
         pretrained_ckpt = os.path.join(model_folder, "checkpoints", cfg.load_checkpoint_name)
         assert cfg.checkpoint_path is None, "You cannot provide both a checkpoint path and a checkpoint run id"
     elif cfg.checkpoint_path is not None and cfg.checkpoint_path != "":
+        model_folder = None
         pretrained_ckpt = cfg.checkpoint_path
     else:
         model_folder, pretrained_ckpt = None, None

configs/model/core/resenc_unet.yaml

@@ -11,6 +11,10 @@ _seg_net:
 _cls_net:
   _target_: asparagus.modules.networks.resenc_unet.${model.cls_net}
   dimensions: ${model.dimensions}
+  dropout_op_kwargs:
+    encoder_dropout_rate: ${model.encoder_dropout_rate}
+    decoder_dropout_rate: ${model.decoder_dropout_rate}
+    inplace: true
 
 _plugin_seg_net:
   _target_: asparagus.modules.networks.resenc_unet.${model.plugin_seg_net}

configs/model/resenc_unet_b.yaml

@@ -6,8 +6,11 @@ dimensions: 3D
 pretrain_net: resenc_unet_b
 
 seg_net: resenc_unet_b
-cls_net: 
+cls_net:
 
-plugin_seg_net: 
+plugin_seg_net:
 plugin_seg_train_n_last_params: 10
 
+# Dropout rates (0.0 = no dropout)
+encoder_dropout_rate: 0.0
+decoder_dropout_rate: 0.0

configs/model/resenc_unet_b_clsreg.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - core/resenc_unet@
+
+dimensions: 3D
+
+pretrain_net:
+
+seg_net:
+cls_net: resenc_unet_b_clsreg
+
+plugin_seg_net:
+plugin_seg_train_n_last_params:
+
+# Dropout rates (0.0 = no dropout)
+encoder_dropout_rate: 0.0
+decoder_dropout_rate: 0.2