Add unet encoder-decoder and image output feature

examples/semantic_segmentation/camseq.py

@@ -0,0 +1,57 @@
+import logging
+import os
+import shutil
+
+import pandas as pd
+import torch
+import yaml
+from torchvision.utils import save_image
+
+from ludwig.api import LudwigModel
+from ludwig.datasets import camseq
+
+# clean out prior results
+shutil.rmtree("./results", ignore_errors=True)
+
+# set up Python dictionary to hold model training parameters
+with open("./config_camseq.yaml") as f:
+    config = yaml.safe_load(f.read())
+
+# Define Ludwig model object that drive model training
+model = LudwigModel(config, logging_level=logging.INFO)
+
+# load Camseq dataset
+df = camseq.load(split=False)
+
+pred_set = df[0:1]  # prediction hold-out 1 image
+data_set = df[1:]  # train,test,validate on remaining images
+
+# initiate model training
+(train_stats, _, output_directory) = model.train(  # training statistics  # location for training results saved to disk
+    dataset=data_set,
+    experiment_name="simple_image_experiment",
+    model_name="single_model",
+    skip_save_processed_input=True,
+)
+
+# print("{}".format(model.model))
+
+# predict
+pred_set.reset_index(inplace=True)
+pred_out_df, results = model.predict(pred_set)
+
+if not isinstance(pred_out_df, pd.DataFrame):
+    pred_out_df = pred_out_df.compute()
+pred_out_df["image_path"] = pred_set["image_path"]
+pred_out_df["mask_path"] = pred_set["mask_path"]
+
+for index, row in pred_out_df.iterrows():
+    pred_mask = torch.from_numpy(row["mask_path_predictions"])
+    pred_mask_path = os.path.dirname(os.path.realpath(__file__)) + "/predicted_" + os.path.basename(row["mask_path"])
+    print(f"\nSaving predicted mask to {pred_mask_path}")
+    if torch.any(pred_mask.gt(1)):
+        pred_mask = pred_mask.float() / 255
+    save_image(pred_mask, pred_mask_path)
+    print("Input image_path:    {}".format(row["image_path"]))
+    print("Label mask_path:     {}".format(row["mask_path"]))
+    print(f"Predicted mask_path: {pred_mask_path}")

examples/semantic_segmentation/config_camseq.yaml

@@ -0,0 +1,33 @@
+input_features:
+  - name: image_path
+    type: image
+    preprocessing:
+      num_processes: 6
+      infer_image_max_height: 1024
+      infer_image_max_width: 1024
+    encoder: unet
+
+output_features:
+  - name: mask_path
+    type: image
+    preprocessing:
+      num_processes: 6
+      infer_image_max_height: 1024
+      infer_image_max_width: 1024
+      infer_image_num_classes: true
+      num_classes: 32
+    decoder:
+      type: unet
+      num_fc_layers: 0
+    loss:
+      type: softmax_cross_entropy
+
+combiner:
+  type: concat
+  num_fc_layers: 0
+
+trainer:
+  epochs: 100
+  early_stop: -1
+  batch_size: 1
+  max_batch_size: 1

ludwig/constants.py

@@ -41,6 +41,8 @@
 INFER_IMAGE_MAX_HEIGHT = "infer_image_max_height"
 INFER_IMAGE_MAX_WIDTH = "infer_image_max_width"
 INFER_IMAGE_SAMPLE_SIZE = "infer_image_sample_size"
+INFER_IMAGE_NUM_CLASSES = "infer_image_num_classes"
+IMAGE_MAX_CLASSES = 128
 NUM_CLASSES = "num_classes"
 NUM_CHANNELS = "num_channels"
 REQUIRES_EQUAL_DIMENSIONS = "requires_equal_dimensions"

ludwig/datasets/configs/camseq.yaml

@@ -0,0 +1,21 @@
+version: 1.0
+name: camseq
+kaggle_dataset_id: carlolepelaars/camseq-semantic-segmentation
+archive_filenames: camseq-semantic-segmentation.zip
+sha256:
+   camseq-semantic-segmentation.zip: ea3aeba2661d9b3e3ea406668e7d9240cb2ba0c7e374914bb6d866147faff502
+loader: camseq.CamseqLoader
+preserve_paths:
+  - images
+  - masks
+description: |
+  CamSeq01 Cambridge Labeled Objects in Video
+  https://www.kaggle.com/datasets/carlolepelaars/camseq-semantic-segmentation
+columns:
+  - name: image_path
+    type: image
+  - name: mask_path
+    type: image
+output_features:
+  - name: mask_path
+    type: image

ludwig/datasets/loaders/camseq.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2023 Aizen Corp.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import os
+from typing import List
+
+import pandas as pd
+
+from ludwig.datasets.loaders.dataset_loader import DatasetLoader
+from ludwig.utils.fs_utils import makedirs
+
+
+class CamseqLoader(DatasetLoader):
+    def transform_files(self, file_paths: List[str]) -> List[str]:
+        if not os.path.exists(self.processed_dataset_dir):
+            os.makedirs(self.processed_dataset_dir)
+
+        # move images and masks into separate directories
+        source_dir = self.raw_dataset_dir
+        images_dir = os.path.join(source_dir, "images")
+        masks_dir = os.path.join(source_dir, "masks")
+        makedirs(images_dir, exist_ok=True)
+        makedirs(masks_dir, exist_ok=True)
+
+        data_files = []
+        for f in os.listdir(source_dir):
+            if f.endswith("_L.png"):  # masks
+                dest_file = os.path.join(masks_dir, f)
+            elif f.endswith(".png"):  # images
+                dest_file = os.path.join(images_dir, f)
+            else:
+                continue
+            source_file = os.path.join(source_dir, f)
+            os.replace(source_file, dest_file)
+            data_files.append(dest_file)
+
+        return super().transform_files(data_files)
+
+    def load_unprocessed_dataframe(self, file_paths: List[str]) -> pd.DataFrame:
+        """Creates a dataframe of image paths and mask paths."""
+        images_dir = os.path.join(self.processed_dataset_dir, "images")
+        masks_dir = os.path.join(self.processed_dataset_dir, "masks")
+        images = []
+        masks = []
+        for f in os.listdir(images_dir):
+            images.append(os.path.join(images_dir, f))
+            mask_f = f[:-4] + "_L.png"
+            masks.append(os.path.join(masks_dir, mask_f))
+
+        return pd.DataFrame({"image_path": images, "mask_path": masks})

ludwig/decoders/__init__.py

@@ -1,5 +1,6 @@
 # register all decoders
 import ludwig.decoders.generic_decoders  # noqa
+import ludwig.decoders.image_decoders  # noqa
 import ludwig.decoders.llm_decoders  # noqa
 import ludwig.decoders.sequence_decoders  # noqa
 import ludwig.decoders.sequence_tagger  # noqa

ludwig/decoders/image_decoders.py

@@ -0,0 +1,91 @@
+#! /usr/bin/env python
+# Copyright (c) 2023 Aizen Corp.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import logging
+from typing import Dict, Optional, Type
+
+import torch
+
+from ludwig.api_annotations import DeveloperAPI
+from ludwig.constants import ENCODER_OUTPUT_STATE, HIDDEN, IMAGE, LOGITS, PREDICTIONS
+from ludwig.decoders.base import Decoder
+from ludwig.decoders.registry import register_decoder
+from ludwig.modules.convolutional_modules import UNetUpStack
+from ludwig.schema.decoders.image_decoders import ImageDecoderConfig, UNetDecoderConfig
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+@register_decoder("unet", IMAGE)
+class UNetDecoder(Decoder):
+    def __init__(
+        self,
+        input_size: int,
+        height: int,
+        width: int,
+        num_channels: int = 1,
+        num_classes: int = 2,
+        conv_norm: Optional[str] = None,
+        decoder_config=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.config = decoder_config
+        self.num_classes = num_classes
+
+        logger.debug(f" {self.name}")
+        if num_classes < 2:
+            raise ValueError(f"Invalid `num_classes` {num_classes} for unet decoder")
+        if height % 16 or width % 16:
+            raise ValueError(f"Invalid `height` {height} or `width` {width} for unet decoder")
+
+        self.unet = UNetUpStack(
+            img_height=height,
+            img_width=width,
+            out_channels=num_classes,
+            norm=conv_norm,
+        )
+
+        self.input_reshape = list(self.unet.input_shape)
+        self.input_reshape.insert(0, -1)
+        self._output_shape = (height, width)
+
+    def forward(self, combiner_outputs: Dict[str, torch.Tensor], target: torch.Tensor):
+        hidden = combiner_outputs[HIDDEN]
+        skips = combiner_outputs[ENCODER_OUTPUT_STATE]
+
+        # unflatten combiner outputs
+        hidden = hidden.reshape(self.input_reshape)
+
+        logits = self.unet(hidden, skips)
+        predictions = logits.argmax(dim=1).squeeze(1).byte()
+
+        return {LOGITS: logits, PREDICTIONS: predictions}
+
+    def get_prediction_set(self):
+        return {LOGITS, PREDICTIONS}
+
+    @staticmethod
+    def get_schema_cls() -> Type[ImageDecoderConfig]:
+        return UNetDecoderConfig
+
+    @property
+    def output_shape(self) -> torch.Size:
+        return torch.Size(self._output_shape)
+
+    @property
+    def input_shape(self) -> torch.Size:
+        return self.unet.input_shape

ludwig/encoders/image/base.py

@@ -19,18 +19,19 @@
 import torch
 
 from ludwig.api_annotations import DeveloperAPI
-from ludwig.constants import ENCODER_OUTPUT, IMAGE
+from ludwig.constants import ENCODER_OUTPUT, ENCODER_OUTPUT_STATE, IMAGE
 from ludwig.encoders.base import Encoder
 from ludwig.encoders.registry import register_encoder
 from ludwig.encoders.types import EncoderOutputDict
-from ludwig.modules.convolutional_modules import Conv2DStack, ResNet
+from ludwig.modules.convolutional_modules import Conv2DStack, ResNet, UNetDownStack
 from ludwig.modules.fully_connected_modules import FCStack
 from ludwig.modules.mlp_mixer_modules import MLPMixer
 from ludwig.schema.encoders.image.base import (
     ImageEncoderConfig,
     MLPMixerConfig,
     ResNetConfig,
     Stacked2DCNNConfig,
+    UNetEncoderConfig,
     ViTConfig,
 )
 from ludwig.utils.torch_utils import FreezeModule
@@ -424,3 +425,46 @@ def input_shape(self) -> torch.Size:
     @property
     def output_shape(self) -> torch.Size:
         return torch.Size(self._output_shape)
+
+
+@DeveloperAPI
+@register_encoder("unet", IMAGE)
+class UNetEncoder(ImageEncoder):
+    def __init__(
+        self,
+        height: int,
+        width: int,
+        num_channels: int = 3,
+        conv_norm: Optional[str] = None,
+        encoder_config=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.config = encoder_config
+
+        logger.debug(f" {self.name}")
+        if height % 16 or width % 16:
+            raise ValueError(f"Invalid `height` {height} or `width` {width} for unet encoder")
+
+        self.unet = UNetDownStack(
+            img_height=height,
+            img_width=width,
+            in_channels=num_channels,
+            norm=conv_norm,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> EncoderOutputDict:
+        hidden, skips = self.unet(inputs)
+        return {ENCODER_OUTPUT: hidden, ENCODER_OUTPUT_STATE: skips}
+
+    @staticmethod
+    def get_schema_cls() -> Type[ImageEncoderConfig]:
+        return UNetEncoderConfig
+
+    @property
+    def output_shape(self) -> torch.Size:
+        return self.unet.output_shape
+
+    @property
+    def input_shape(self) -> torch.Size:
+        return self.unet.input_shape

ludwig/features/feature_registries.py

@@ -43,7 +43,7 @@
 )
 from ludwig.features.date_feature import DateFeatureMixin, DateInputFeature
 from ludwig.features.h3_feature import H3FeatureMixin, H3InputFeature
-from ludwig.features.image_feature import ImageFeatureMixin, ImageInputFeature
+from ludwig.features.image_feature import ImageFeatureMixin, ImageInputFeature, ImageOutputFeature
 from ludwig.features.number_feature import NumberFeatureMixin, NumberInputFeature, NumberOutputFeature
 from ludwig.features.sequence_feature import SequenceFeatureMixin, SequenceInputFeature, SequenceOutputFeature
 from ludwig.features.set_feature import SetFeatureMixin, SetInputFeature, SetOutputFeature
@@ -108,6 +108,7 @@ def get_output_type_registry() -> Dict:
         TIMESERIES: TimeseriesOutputFeature,
         VECTOR: VectorOutputFeature,
         CATEGORY_DISTRIBUTION: CategoryDistributionOutputFeature,
+        IMAGE: ImageOutputFeature,
     }