Add metrics (qubvel-org#531)

* Add metrics 
* Add docs
* Add example notebook

Author: qubvel

README.md

@@ -12,8 +12,9 @@ The main features of this library are:
 
  - High level API (just two lines to create a neural network)
  - 9 models architectures for binary and multi class segmentation (including legendary Unet)
- - 113 available encoders
+ - 113 available encoders (and 400+ encoders from [timm](https://github.com/rwightman/pytorch-image-models))
  - All encoders have pre-trained weights for faster and better convergence
+ - Popular metrics and losses for training routines
 
 ### [📚 Project Documentation 📚](http://smp.readthedocs.io/)
 
@@ -68,9 +69,10 @@ preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')
 Congratulations! You are done! Now you can train your model with your favorite framework!
 
 ### 💡 Examples <a name="examples"></a>
+ - Training model for pets binary segmentation with Pytorch-Lightning [notebook](https://github.com/qubvel/segmentation_models.pytorch/blob/master/examples/binary_segmentation_intro.ipynb) and [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googlecolab/colabtools/blob/master/examples/binary_segmentation_intro.ipynb)
  - Training model for cars segmentation on CamVid dataset [here](https://github.com/qubvel/segmentation_models.pytorch/blob/master/examples/cars%20segmentation%20(camvid).ipynb).
- - Training SMP model with [Catalyst](https://github.com/catalyst-team/catalyst) (high-level framework for PyTorch), [TTAch](https://github.com/qubvel/ttach) (TTA library for PyTorch) and [Albumentations](https://github.com/albu/albumentations) (fast image augmentation library) - [here](https://github.com/catalyst-team/catalyst/blob/master/examples/notebooks/segmentation-tutorial.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/catalyst-team/catalyst/blob/master/examples/notebooks/segmentation-tutorial.ipynb)
- - Training SMP model with [Pytorch-Lightning](https://pytorch-lightning.readthedocs.io) framework - [here](https://github.com/ternaus/cloths_segmentation) (clothes binary segmentation by [@teranus](https://github.com/ternaus)).
+ - Training SMP model with [Catalyst](https://github.com/catalyst-team/catalyst) (high-level framework for PyTorch), [TTAch](https://github.com/qubvel/ttach) (TTA library for PyTorch) and [Albumentations](https://github.com/albu/albumentations) (fast image augmentation library) - [here](https://github.com/catalyst-team/catalyst/blob/v21.02rc0/examples/notebooks/segmentation-tutorial.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/catalyst-team/catalyst/blob/v21.02rc0/examples/notebooks/segmentation-tutorial.ipynb)
+ - Training SMP model with [Pytorch-Lightning](https://pytorch-lightning.readthedocs.io) framework - [here](https://github.com/ternaus/cloths_segmentation) (clothes binary segmentation by [@ternaus](https://github.com/ternaus)).
 
 ### 📦 Models <a name="models"></a>
 

docs/conf.py

@@ -46,6 +46,7 @@ def get_version():
     'sphinx.ext.napoleon',
     'sphinx.ext.viewcode',
     'sphinx.ext.mathjax', 
+    'autodocsumm',
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -95,6 +96,8 @@ def get_version():
     'tqdm',
     'numpy',
     'timm',
+    'cv2',
+    'PIL',
     'pretrainedmodels',
     'torchvision',
     'efficientnet-pytorch',
@@ -118,3 +121,8 @@ def f(app, obj, bound_method):
 
 def setup(app):
     app.connect('autodoc-before-process-signature', f)
+
+
+# Custom configuration --------------------------------------------------------
+
+autodoc_member_order = 'bysource'

docs/index.rst

@@ -16,6 +16,7 @@ Welcome to Segmentation Models's documentation!
    encoders
    encoders_timm
    losses
+   metrics
    insights
 
 

docs/metrics.rst

@@ -0,0 +1,8 @@
+📈 Metrics
+==========
+
+Functional metrics
+~~~~~~~~~~~~~~~~~~
+.. automodule:: segmentation_models_pytorch.metrics.functional
+        :members:
+        :autosummary:

docs/requirements.txt

@@ -1,2 +1,3 @@
 faculty-sphinx-theme==0.2.2
-six==1.15.0
+six==1.15.0
+autodocsumm

examples/binary_segmentation_intro.ipynb

@@ -4,4 +4,4 @@ efficientnet-pytorch==0.6.3
 timm==0.4.12
 
 tqdm
-opencv-python-headless
+pillow

requirements.txt

@@ -1,3 +1,9 @@
+from . import datasets
+from . import encoders
+from . import decoders
+from . import losses
+from . import metrics
+
 from .decoders.unet import Unet
 from .decoders.unetplusplus import UnetPlusPlus
 from .decoders.manet import MAnet
@@ -7,27 +13,23 @@
 from .decoders.deeplabv3 import DeepLabV3, DeepLabV3Plus
 from .decoders.pan import PAN
 
-from . import encoders
-from . import decoders
-from . import losses
-
 from .__version__ import __version__
 
-from typing import Optional
-import torch
+# some private imports for create_model function
+from typing import Optional as _Optional
+import torch as _torch
 
 
 def create_model(
     arch: str,
     encoder_name: str = "resnet34",
-    encoder_weights: Optional[str] = "imagenet",
+    encoder_weights: _Optional[str] = "imagenet",
     in_channels: int = 3,
     classes: int = 1,
     **kwargs,
-) -> torch.nn.Module:
-    """Models wrapper. Allows to create any model just with parametes
-
-    """
+) -> _torch.nn.Module:
+    """Models entrypoint, allows to create any model architecture just with 
+    parameters, without using its class"""
 
     archs = [Unet, UnetPlusPlus, MAnet, Linknet, FPN, PSPNet, DeepLabV3, DeepLabV3Plus, PAN]
     archs_dict = {a.__name__.lower(): a for a in archs}

segmentation_models_pytorch/__init__.py

@@ -1,7 +1,9 @@
 import os
-import cv2
+import torch
 import shutil
 import numpy as np
+
+from PIL import Image
 from tqdm import tqdm
 from urllib.request import urlretrieve
 
@@ -15,8 +17,6 @@ def __init__(self, root, mode="train", transform=None):
         self.root = root
         self.mode = mode
         self.transform = transform
-        
-        self._download_dataset() # download only if it does not exist
 
         self.images_directory = os.path.join(self.root, "images")
         self.masks_directory = os.path.join(self.root, "annotations", "trimaps")
@@ -32,10 +32,9 @@ def __getitem__(self, idx):
         image_path = os.path.join(self.images_directory, filename + ".jpg")
         mask_path = os.path.join(self.masks_directory, filename + ".png")
 
-        image = cv2.imread(image_path)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = np.array(Image.open(image_path).convert("RGB"))
 
-        trimap = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
+        trimap = np.array(Image.open(mask_path))
         mask = self._preprocess_mask(trimap)
 
         sample = dict(image=image, mask=mask, trimap=trimap)
@@ -63,34 +62,33 @@ def _read_split(self):
             filenames = [x for i, x in enumerate(filenames) if i % 10 == 0]
         return filenames
 
-    def _download_dataset(self):
+    @staticmethod
+    def download(root):
 
         # load images
-        filepath = os.path.join(self.root, "images.tar.gz")
+        filepath = os.path.join(root, "images.tar.gz")
         download_url(
             url="https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz", filepath=filepath,
         )
         extract_archive(filepath)
 
         # load annotations
-        filepath = os.path.join(self.root, "annotations.tar.gz")
+        filepath = os.path.join(root, "annotations.tar.gz")
         download_url(
             url="https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz", filepath=filepath,
         )
         extract_archive(filepath)
 
 
 class SimpleOxfordPetDataset(OxfordPetDataset):
-    """Dataset for example without augmentations and transforms"""
-
     def __getitem__(self, *args, **kwargs):
 
         sample = super().__getitem__(*args, **kwargs)
 
         # resize images
-        image = cv2.resize(sample["image"], (256, 256), cv2.INTER_LINEAR)
-        mask = cv2.resize(sample["mask"], (256, 256), cv2.INTER_NEAREST)
-        trimap = cv2.resize(sample["trimap"], (256, 256), cv2.INTER_NEAREST)
+        image = np.array(Image.fromarray(sample["image"]).resize((256, 256), Image.LINEAR))
+        mask = np.array(Image.fromarray(sample["mask"]).resize((256, 256), Image.NEAREST))
+        trimap = np.array(Image.fromarray(sample["trimap"]).resize((256, 256), Image.NEAREST))
 
         # convert to other format HWC -> CHW
         sample["image"] = np.moveaxis(image, -1, 0)

segmentation_models_pytorch/datasets/oxford_pet.py

@@ -0,0 +1,20 @@
+from .functional import (
+    get_stats,
+    fbeta_score,
+    f1_score,
+    iou_score,
+    accuracy,
+    precision,
+    recall,
+    sensitivity,
+    specificity,
+    balanced_accuracy,
+    positive_predictive_value,
+    negative_predictive_value,
+    false_negative_rate,
+    false_positive_rate,
+    false_discovery_rate,
+    false_omission_rate,
+    positive_likelihood_ratio,
+    negative_likelihood_ratio,
+)