Initial mmsegmentation support (airctic#934)

* Fixed PIL size bug in ImageRecordComponent (airctic#889)

* Preliminary work on mmsegmentation integration

This is currently very crude, mostly copy/pasting of the mmdet integration with a few changes.

* Updated mmseg dataloaders based on unet implementation

* Training now kind of working!

- Fixed issue with number of classes
- Fixed dataloader issue (mask tensor size)

* Improved predictions handling

* Fixed prediction, general code clean-up

* Simplification of folder structure

* Getting closer to final structure

* Finished first structure update

- This implementation was created based on the existing mmdet integration, and some remaining, unused folders, were removed

* Attempt to fix show_batch for mmseg

* Added binary and multi-class dice coefficient metrics

* Removed test-only code in binary Dice coefficient method

* Initial support for multiple pre-trained variants

* Reformatting using black

* First version of DeepLabV3plus support

* Exceptions to elegantly handled unknown pre-trained variants

* Added all DeepLabV3 pretrained variants

* Ground truth masks saved with predictions if keep_images set to True

* Added all DeepLabV3Plus pre-training variants

* Added proper support for param_groups

* Removed erroneous pre-trained variants for DeepLabV3Plus

* Fixed erroneous DeepLabV3 pre-trained variants

* re-added default values to DeepLabV3 pretrained variants

* Improved model loading and initialization

* Improved how distributed BNs are handled

* - Proper integration of loop_mmseg
- First test!

* Updated tests

* __init__.py file

* jaccard index metric

* update init file to include multilabel dice coef

* updates to logical statements

* update to handle denominator equal to 0

* update to handle denominator equal to 0

* removed repeated code

* removed repeated code

* Getting started with seg notebook

* Formatting fix (black)

* Removed temporary file

* Added mmsegmentation to Dockerfile

* Added mmseg installation to workflows

* Updating mmcv to 1.3.7 for mmseg support

* Added artifacts to gitignore (wandb)

* Testing mim-based install in actions

* Fixing mim-based install

* Pinning mmcv version in mim installs

* Bumping mmcv-full to 1.3.13

* Improved CPU support

* Reverted workflow files to match current master

* Improved tests on CPU devices

* Update docs action to match new dependencies versions

* Better handling of the device mess?

* Attempt to remove hacky device code

* Added mmseg to soft dependency test

* changed to binary jaccard index

* delete jaccard_index

* added jaccard index

* up to date metric test files

* Fixed init for binary jaccard index

* Argument to exclude classes from multiclass dice coefficient metric

* Added background exclusion case for multilabel dice coefficient tests

* adjusted setup, need to verify values

* added comment for correct values

* updated cases and setup

* updated cases and setup

* added class map to dictionary for each case

* Adapted RLE mask saving code, might need to revisit later

* Added support for wandb logging of semantic seg masks

* Added option not to pad dimensions when creating a MaskArray object

* Fixed typo

* Resampling of masks should use nearest neighbour

* Added TODO to masks.py

* Fixed loss keys for mssg

* mmseg installation in icevision_install.sh

* Fixed typo in install script

* Fixed albumentation for polygon mask

* Black formatting

* Fix handlig of special case of Polygon masks

* Started updating getting started notebook

* More updates to the notebook

* Adding after_pred convert raw pred callback to unet fastai learner

* Fixed unet prediction when GT not available

* More fixes to unet prediction callback

* Fixing predict batch for mmseg models

* Fixed predictions conversion for mmseg

* Improved unet prediction conversion

* Black formatting

* segformer support

* Updated mmseg model creation function

* Misc comment

* Added options to install script to assist with testing (temporary)

* Black formatting

* Remove install colab script added by mistake

* Actual updated install script

* Updated semantic seg notebook

* Removed legacy cell from notebook

* Further updated notebook

* Updated "open in colab" icon URL

* Added cell to restart kernel after installation

* Reverted notebook changes for merge into master

* Started work on README

* Updated README somewhat

* Updated semantic seg notebook to match updated version on master

* Fix draw data for mask files

* Implementation of mask resizing for MaskFile objects

* Support for MaskFile obnects in draw_data

* Fixed mask display in show_batch for mask-enabled mmdet models

* import typo

* Adding warning if re-creating masks from file in draw_data

* Added warning if maskfile loaded from disc for display

* README typo fixed

* Removed temporary mmsegmentation installation option

* Reverted changes to plot_top_losses

* Disabling mmseg PL implementation for now (pending update)

* Re-added unet callback

* Fixed CI

* Actually disabled PL implementation

* Removing test file uploaded by mistake

* mmseg fastai training test

* Fixed BackboneConfig location

* Fixed formatting

* Use latest convert_raw_predictions for Unet

* Fixed Unet PL test with new raw prediction conversion method

* Removed lightning mmseg files pending proper implementation

Co-authored-by: Gabriella Lanouette <[email protected]>

Author: nicjac

.github/workflows/ci-all-testing.yml

@@ -2,9 +2,9 @@ name: Run unit tests
 
 on:
   push:
-    branches: [ master ]
+    branches: [master]
   pull_request:
-    branches: [ master ]
+    branches: [master]
 
 jobs:
   test:
@@ -16,33 +16,32 @@ jobs:
         python-version: [3.7, 3.8]
 
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python-version }}
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
 
-    - name: Install package
-      run: |
-        sh ./icevision_install.sh cpu
-        pip install -e ".[all,dev]"
-        pip install fiftyone
+      - name: Install package
+        run: |
+          sh ./icevision_install.sh cpu
+          pip install -e ".[all,dev]"
+          pip install fiftyone
 
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+          flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+      - name: Unit tests
+        run: pytest tests -m "not cuda" --cov=icevision --cov-report=xml --color=yes
 
-    - name: Unit tests
-      run: pytest tests -m "not cuda" --cov=icevision --cov-report=xml --color=yes
-
-    - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v1
-      with:
-        token: ${{ secrets.CODECOV_TOKEN }}
-        file: ./coverage.xml
-        flags: unittests
-        name: codecov-umbrella
-        fail_ci_if_error: false
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v1
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: ./coverage.xml
+          flags: unittests
+          name: codecov-umbrella
+          fail_ci_if_error: false

.github/workflows/mk-docs-build.yml

@@ -2,18 +2,17 @@ name: Build mkdocs
 
 on:
   pull_request:
-    branches: [ master ]
+    branches: [master]
   push:
-    branches: [ master ]
+    branches: [master]
   release:
-    types: [ created ]    
+    types: [created]
 
 env:
   SITE_BRANCH: ${{ 'gh-pages-ver' }}
 
 jobs:
   build-docs:
-
     runs-on: ubuntu-18.04
     steps:
       - uses: actions/checkout@v2
@@ -32,35 +31,30 @@ jobs:
           
 
       - name: Prepare the docs
-        run:  |
+        run: |
           cd docs
           python autogen.py
-
       - name: Setup git config
         run: |
           git config user.name "GitHub Actions Bot"
           git config user.email "[email protected]"
-
       - name: Build the docs locally only
         if: github.event_name == 'pull_request'
-        run:  |
+        run: |
           cd docs
-          mike deploy dev -b ${{ env.SITE_BRANCH }} 
-
+          mike deploy dev -b ${{ env.SITE_BRANCH }}
       - name: Deploy dev docs
         id: deploy_dev
         if: github.repository == 'airctic/icevision' && github.event_name == 'push'
         run: |
           cd docs
           mike deploy dev -b ${{ env.SITE_BRANCH }} -p
-          echo '::set-output name=MIKE_VERSIONS::'$(mike list -b ${{ env.SITE_BRANCH }} | wc -l)     
-
+          echo '::set-output name=MIKE_VERSIONS::'$(mike list -b ${{ env.SITE_BRANCH }} | wc -l)
       - name: Set dev as default
         if: steps.deploy_dev.outputs.MIKE_VERSIONS == 1
         run: |
           cd docs
           mike set-default -b ${{ env.SITE_BRANCH }} dev -p
-
       - name: Get latest release tag
         if: github.event_name == 'release' && !github.event.release.prerelease
         id: latest
@@ -75,4 +69,4 @@ jobs:
           cd docs
           echo Deploy as ${{ steps.latest.outputs.release }} 
           mike deploy -b ${{ env.SITE_BRANCH }} ${{ steps.latest.outputs.release }} -p
-          mike set-default -b ${{ env.SITE_BRANCH }} ${{ steps.latest.outputs.release }} -p     
+          mike set-default -b ${{ env.SITE_BRANCH }} ${{ steps.latest.outputs.release }} -p

.gitignore

@@ -154,4 +154,10 @@ archives/
 *.pth
 notebooks/wandb/latest-run
 
+artifacts/*
 checkpoints
+
+wandb
+wandb/*
+
+*/artifacts/*

docker/Dockerfile

@@ -6,4 +6,5 @@ RUN pip install git+https://github.com/airctic/icedata.git -U
 RUN pip install yolov5-icevision -U 
 RUN pip install mmcv-full==1.3.7 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.8.0/index.html -U
 RUN pip install mmdet==2.13.0 -U
+RUN pip install mmsegmentation==0.17.0 -U
 RUN pip install ipywidgets

icevision/core/exceptions.py

@@ -1,4 +1,10 @@
-__all__ = ["InvalidDataError", "AutofixAbort", "AbortParseRecord"]
+__all__ = [
+    "InvalidDataError",
+    "AutofixAbort",
+    "AbortParseRecord",
+    "InvalidMMSegModelType",
+    "PreTrainedVariantNotFound",
+]
 
 
 class InvalidDataError(Exception):
@@ -11,3 +17,11 @@ class AutofixAbort(Exception):
 
 class AbortParseRecord(Exception):
     pass
+
+
+class InvalidMMSegModelType(Exception):
+    pass
+
+
+class PreTrainedVariantNotFound(Exception):
+    pass

icevision/core/mask.py

@@ -64,10 +64,11 @@ class MaskArray(Mask):
 
     # Arguments
         data: Mask array, with the dimensions: (num_instances, height, width)
+        pad_dim: bool
     """
 
-    def __init__(self, data: np.uint8):
-        if len(data.shape) == 2:
+    def __init__(self, data: np.uint8, pad_dim: bool = True):
+        if pad_dim and (len(data.shape) == 2):
             data = np.expand_dims(data, 0)
         self.data = data.astype(np.uint8)
 
@@ -137,10 +138,21 @@ class MaskFile(Mask):
     def __init__(self, filepath: Union[str, Path]):
         self.filepath = Path(filepath)
 
-    def to_mask(self, h, w):
-        mask = np.array(open_img(self.filepath, gray=True))
+    def to_mask(self, h=None, w=None):
+        mask_img = open_img(self.filepath, gray=True)
+
+        if (h is not None) and (w is not None):
+            # If the dimensions provided in h and w do not match the size of the mask, resize the mask accordingly
+            (w_org, h_org) = mask_img.size
+
+            # TODO: Check NEAREST is always the best option or only for binary?
+            if w_org != w or h_org != h:
+                mask_img = mask_img.resize((w, h), resample=PIL.Image.NEAREST)
+
+        mask = np.array(mask_img)
         obj_ids = np.unique(mask)[1:]
         masks = mask == obj_ids[:, None, None]
+
         return MaskArray(masks)
 
     def to_coco_rle(self, h, w) -> List[dict]:
@@ -275,17 +287,26 @@ def __init__(self, filepath: Union[str, Path], binary=False):
         self.filepath = Path(filepath)
         self.binary = binary
 
-    def to_mask(self, h, w):
+    def to_mask(self, h, w, pad_dim=True):
         # TODO: convert the 255 masks
         mask = open_img(self.filepath, gray=True)
+
+        # If the dimensions provided in h and w do not match the size of the mask, resize the mask accordingly
+        (w_org, h_org) = mask.size
+
+        # TODO: Check NEAREST is always the best option or only for binary?
+        if w_org != w or h_org != h:
+            mask = mask.resize((w, h), resample=PIL.Image.NEAREST)
+
         # HACK: because open_img now return PIL
         mask = np.array(mask)
 
         # convert 255 pixels to 1
         if self.binary:
             mask[mask == 255] = 1
 
-        return MaskArray(mask[None])
+        # control array padding behaviour
+        return MaskArray(mask, pad_dim=pad_dim)
 
     def to_coco_rle(self, h, w) -> List[dict]:
         raise NotImplementedError

icevision/metrics/__init__.py

@@ -1,3 +1,6 @@
 from icevision.metrics.metric import *
 from icevision.metrics.coco_metric import *
 from icevision.metrics.confusion_matrix import *
+from icevision.metrics.segmentation_accuracy import *
+from icevision.metrics.dice_coefficient import *
+from icevision.metrics.jaccard_index import *

icevision/metrics/dice_coefficient/__init__.py

@@ -0,0 +1,2 @@
+from icevision.metrics.dice_coefficient.binary_dice_coefficient import *
+from icevision.metrics.dice_coefficient.multiclass_dice_coefficient import *

icevision/metrics/dice_coefficient/binary_dice_coefficient.py

@@ -0,0 +1,49 @@
+__all__ = ["BinaryDiceCoefficient"]
+
+from icevision.imports import *
+from icevision.utils import *
+from icevision.data import *
+from icevision.metrics.metric import *
+
+
+class BinaryDiceCoefficient(Metric):
+    """Binary Dice Coefficient for Semantic Segmentation
+
+    Calculates Dice Coefficient for semantic segmentation (binary images only).
+
+    """
+
+    def __init__(self):
+        self._reset()
+
+    def _reset(self):
+        self._union = 0
+        self._intersection = 0
+
+    def accumulate(self, preds):
+
+        pred = (
+            np.stack([x.pred.segmentation.mask_array.data for x in preds])
+            .astype(np.bool)
+            .flatten()
+        )
+
+        target = (
+            np.stack([x.ground_truth.segmentation.mask_array.data for x in preds])
+            .astype(np.bool)
+            .flatten()
+        )
+
+        self._union += pred.sum() + target.sum()
+        self._intersection += np.logical_and(pred, target).sum()
+
+    def finalize(self) -> Dict[str, float]:
+
+        if self._union == 0:
+            dice = 0
+
+        else:
+            dice = 2.0 * self._intersection / self._union
+
+        self._reset()
+        return {"dummy_value_for_fastai": dice}

icevision/metrics/dice_coefficient/multiclass_dice_coefficient.py

@@ -0,0 +1,73 @@
+__all__ = ["MulticlassDiceCoefficient"]
+
+from icevision.imports import *
+from icevision.utils import *
+from icevision.data import *
+from icevision.metrics.metric import *
+
+
+class MulticlassDiceCoefficient(Metric):
+    """Multi-class Dice Coefficient for Semantic Segmentation
+
+    Calculates Dice Coefficient for semantic segmentation (multi-class tasks).
+
+    Heaviliy inspired by fastai's implementation:  Implement multi-class version following https://github.com/fastai/fastai/blob/594e1cc20068b0d99bfc30bfe6dac88ab381a157/fastai/metrics.py#L343
+
+    # Arguments
+        classes_to_exclude: A list of class names to exclude from metric computation
+    """
+
+    def __init__(self, classes_to_exclude: list = []):
+        self.classes_to_exclude = classes_to_exclude
+        self._reset()
+
+    def _reset(self):
+        self._union = {}
+        self._intersection = {}
+
+    def accumulate(self, preds):
+
+        pred = np.stack([x.pred.segmentation.mask_array.data for x in preds]).flatten()
+
+        target = self._seg_masks_gt = np.stack(
+            [x.ground_truth.segmentation.mask_array.data for x in preds]
+        ).flatten()
+
+        unique_classes_id = list(preds[0].segmentation.class_map._class2id.values())
+
+        for c in self.classes_to_exclude:
+            if c in preds[0].segmentation.class_map._class2id:
+                unique_classes_id.remove(preds[0].segmentation.class_map._class2id[c])
+
+        # We iterate through all unique classes across preditions and targets
+        for c in unique_classes_id:
+
+            p = np.where(pred == c, 1, 0)
+            t = np.where(target == c, 1, 0)
+
+            c_inter = np.logical_and(p, t).sum()
+            c_union = p.sum() + t.sum()
+
+            if c in self._intersection:
+                self._intersection[c] += c_inter
+                self._union[c] += c_union
+            else:
+                self._intersection[c] = c_inter
+                self._union[c] = c_union
+
+    def finalize(self) -> Dict[str, float]:
+
+        class_dice_scores = np.array([])
+
+        for c in self._intersection:
+            class_dice_scores = np.append(
+                class_dice_scores,
+                2.0 * self._intersection[c] / self._union[c]
+                if self._union[c] > 0
+                else np.nan,
+            )
+
+        dice = np.nanmean(class_dice_scores)
+
+        self._reset()
+        return {"dummy_value_for_fastai": dice}

icevision/metrics/jaccard_index/__init__.py

@@ -0,0 +1 @@
+from icevision.metrics.jaccard_index.binary_jaccard_index import *