Full Code of qubvel-org/segmentation_models.pytorch for AI

Repository: qubvel-org/segmentation_models.pytorch
Branch: main
Commit: 91c0895ef8d8
Files: 174
Total size: 6.3 MB

Directory structure:
gitextract_j1x_q930/

├── .devcontainer/
│   └── devcontainer.json
├── .github/
│   ├── FUNDING.yml
│   ├── codecov.yml
│   ├── dependabot.yml
│   └── workflows/
│       ├── pypi.yml
│       ├── stale.yaml
│       └── tests.yml
├── .gitignore
├── .readthedocs.yaml
├── HALLOFFAME.md
├── LICENSE
├── Makefile
├── README.md
├── docs/
│   ├── Makefile
│   ├── conf.py
│   ├── encoders.rst
│   ├── encoders_dpt.rst
│   ├── encoders_timm.rst
│   ├── index.rst
│   ├── insights.rst
│   ├── install.rst
│   ├── losses.rst
│   ├── make.bat
│   ├── metrics.rst
│   ├── models.rst
│   ├── quickstart.rst
│   └── save_load.rst
├── examples/
│   ├── binary_segmentation_buildings.py
│   ├── binary_segmentation_intro.ipynb
│   ├── camvid_segmentation_multiclass.ipynb
│   ├── cars segmentation (camvid).ipynb
│   ├── convert_to_onnx.ipynb
│   ├── dpt_inference_pretrained.ipynb
│   ├── save_load_model_and_share_with_hf_hub.ipynb
│   ├── segformer_inference_pretrained.ipynb
│   └── upernet_inference_pretrained.ipynb
├── licenses/
│   ├── LICENSES.md
│   ├── LICENSE_apache.md
│   ├── LICENSE_apple.md
│   └── LICENSE_nvidia.md
├── misc/
│   ├── generate_table.py
│   ├── generate_table_timm.py
│   └── generate_test_models.py
├── pyproject.toml
├── requirements/
│   ├── docs.txt
│   ├── minimum.old
│   ├── required.txt
│   └── test.txt
├── scripts/
│   └── models-conversions/
│       ├── dpt-original-to-smp.py
│       ├── segformer-original-decoder-to-smp.py
│       └── upernet-hf-to-smp.py
├── segmentation_models_pytorch/
│   ├── __init__.py
│   ├── __version__.py
│   ├── base/
│   │   ├── __init__.py
│   │   ├── heads.py
│   │   ├── hub_mixin.py
│   │   ├── initialization.py
│   │   ├── model.py
│   │   ├── modules.py
│   │   └── utils.py
│   ├── datasets/
│   │   ├── __init__.py
│   │   └── oxford_pet.py
│   ├── decoders/
│   │   ├── __init__.py
│   │   ├── deeplabv3/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── dpt/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── fpn/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── linknet/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── manet/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── pan/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── pspnet/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── segformer/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── unet/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   ├── unetplusplus/
│   │   │   ├── __init__.py
│   │   │   ├── decoder.py
│   │   │   └── model.py
│   │   └── upernet/
│   │       ├── __init__.py
│   │       ├── decoder.py
│   │       └── model.py
│   ├── encoders/
│   │   ├── __init__.py
│   │   ├── _base.py
│   │   ├── _dpn.py
│   │   ├── _efficientnet.py
│   │   ├── _inceptionresnetv2.py
│   │   ├── _inceptionv4.py
│   │   ├── _legacy_pretrained_settings.py
│   │   ├── _preprocessing.py
│   │   ├── _senet.py
│   │   ├── _utils.py
│   │   ├── _xception.py
│   │   ├── densenet.py
│   │   ├── dpn.py
│   │   ├── efficientnet.py
│   │   ├── inceptionresnetv2.py
│   │   ├── inceptionv4.py
│   │   ├── mix_transformer.py
│   │   ├── mobilenet.py
│   │   ├── mobileone.py
│   │   ├── resnet.py
│   │   ├── senet.py
│   │   ├── timm_efficientnet.py
│   │   ├── timm_sknet.py
│   │   ├── timm_universal.py
│   │   ├── timm_vit.py
│   │   ├── vgg.py
│   │   └── xception.py
│   ├── losses/
│   │   ├── __init__.py
│   │   ├── _functional.py
│   │   ├── constants.py
│   │   ├── dice.py
│   │   ├── focal.py
│   │   ├── jaccard.py
│   │   ├── lovasz.py
│   │   ├── mcc.py
│   │   ├── soft_bce.py
│   │   ├── soft_ce.py
│   │   └── tversky.py
│   ├── metrics/
│   │   ├── __init__.py
│   │   └── functional.py
│   └── utils/
│       ├── __init__.py
│       ├── base.py
│       ├── functional.py
│       ├── losses.py
│       ├── meter.py
│       ├── metrics.py
│       └── train.py
└── tests/
    ├── __init__.py
    ├── base/
    │   ├── test_freeze_encoder.py
    │   └── test_modules.py
    ├── conftest.py
    ├── encoders/
    │   ├── __init__.py
    │   ├── base.py
    │   ├── test_batchnorm_deprecation.py
    │   ├── test_common.py
    │   ├── test_pretrainedmodels_encoders.py
    │   ├── test_smp_encoders.py
    │   ├── test_timm_ported_encoders.py
    │   ├── test_timm_universal.py
    │   ├── test_timm_vit_encoders.py
    │   └── test_torchvision_encoders.py
    ├── models/
    │   ├── __init__.py
    │   ├── base.py
    │   ├── test_deeplab.py
    │   ├── test_dpt.py
    │   ├── test_fpn.py
    │   ├── test_linknet.py
    │   ├── test_manet.py
    │   ├── test_pan.py
    │   ├── test_psp.py
    │   ├── test_segformer.py
    │   ├── test_unet.py
    │   ├── test_unetplusplus.py
    │   └── test_upernet.py
    ├── test_base.py
    ├── test_losses.py
    ├── test_preprocessing.py
    └── utils.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .devcontainer/devcontainer.json
================================================
{
  "image": "mcr.microsoft.com/devcontainers/universal:2",
  "features": {
    "ghcr.io/devcontainers/features/python:1": {
      "verison": 3.6
    }
  }
}


================================================
FILE: .github/FUNDING.yml
================================================
# These are supported funding model platforms

github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
patreon: # Replace with a single Patreon username
open_collective: # Replace with a single Open Collective username
ko_fi: qubvel
tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
liberapay: qubvel
issuehunt: # Replace with a single IssueHunt username
otechie: # Replace with a single Otechie username
custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']


================================================
FILE: .github/codecov.yml
================================================
comment: false
github_checks:
  annotations: false


================================================
FILE: .github/dependabot.yml
================================================
version: 2
updates:
  - package-ecosystem: "github-actions"
    directory: "/"
    schedule:
      interval: "daily"
  - package-ecosystem: "pip"
    directory: "/requirements"
    schedule:
      interval: "daily"
    groups:
      torch:
        patterns:
          - "torch"
          - "torchvision"
    ignore:
      - dependency-name: "setuptools"
        update-types: ["version-update:semver-patch"]


================================================
FILE: .github/workflows/pypi.yml
================================================
name: Upload Python Package

on:
  release:
    types: [published]

jobs:
  build:
    name: build
    runs-on: ubuntu-latest
    steps:
      - name: Clone repo
        uses: actions/checkout@v6

      - name: Set up python
        uses: actions/setup-python@v6.2.0
        with:
          python-version: '3.14'

      - name: Install pip dependencies
        run: pip install build

      - name: List pip dependencies
        run: pip list

      - name: Build project
        run: python3 -m build

      - name: Upload artifacts
        uses: actions/upload-artifact@v7.0.0
        with:
          name: pypi-dist
          path: dist/

  pypi:
    name: pypi
    needs:
      - build
    permissions:
      id-token: write
    runs-on: ubuntu-latest
    steps:
      - name: Download artifacts
        uses: actions/download-artifact@v8.0.1
        with:
          name: pypi-dist
          path: dist/

      - name: Publish to PyPI
        uses: pypa/gh-action-pypi-publish@v1.13.0


================================================
FILE: .github/workflows/stale.yaml
================================================
name: 'Close stale issues and PRs'
on:
  schedule:
    - cron: '30 1 * * *'

jobs:
  stale:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/stale@v10
        with:
          stale-issue-message: 'This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.'
          stale-pr-message: 'This PR is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 15 days.'
          close-issue-message: 'This issue was closed because it has been stalled for 7 days with no activity.'
          close-pr-message: 'This PR was closed because it has been stalled for 15 days with no activity.'
          days-before-issue-stale: 60
          days-before-pr-stale: 60
          days-before-issue-close: 7
          days-before-pr-close: 15
          operations-per-run: 100
          ascending: true


================================================
FILE: .github/workflows/tests.yml
================================================

# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions

name: CI

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

jobs:

  style:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6
      - name: Set up Python
        uses: astral-sh/setup-uv@v7
        with:
          python-version: "3.14"
          activate-environment: true

      - name: Install dependencies
        run: uv pip install -r requirements/test.txt

      # Update output format to enable automatic inline annotations.
      - name: Run Ruff Linter
        run: ruff check --output-format=github
      - name: Run Ruff Formatter
        run: ruff format --check

  test:
    strategy:
      matrix:
        os: [ubuntu-latest, macos-latest, windows-latest]
        python-version: ["3.12", "3.13", "3.14"]
    runs-on: ${{ matrix.os }}
    steps:
    - uses: actions/checkout@v6

    - name: Set up Python ${{ matrix.python-version }}
      uses: astral-sh/setup-uv@v7
      with:
        python-version: ${{ matrix.python-version }}
        activate-environment: true

    - name: Install dependencies
      run: uv pip install -r requirements/required.txt -r requirements/test.txt

    - name: Show installed packages
      run: uv pip list

    - name: Test with PyTest
      run: uv run pytest -v -rsx -n 2 --cov=segmentation_models_pytorch --cov-report=xml --cov-config=pyproject.toml --non-marked-only

    - name: Upload coverage reports to Codecov
      uses: codecov/codecov-action@v5
      with:
        token: ${{ secrets.CODECOV_TOKEN }}
        slug: qubvel-org/segmentation_models.pytorch

  test_logits_match:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
      uses: astral-sh/setup-uv@v7
      with:
        python-version: "3.14"
        activate-environment: true
    - name: Install dependencies
      run: uv pip install -r requirements/required.txt -r requirements/test.txt
    - name: Show installed packages
      run: uv pip list
    - name: Test with PyTest
      run: RUN_SLOW=1 uv run pytest -v -rsx -n 2 -m "logits_match"

  test_torch_compile:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
      uses: astral-sh/setup-uv@v7
      with:
        python-version: "3.14"
        activate-environment: true
    - name: Install dependencies
      run: uv pip install -r requirements/required.txt -r requirements/test.txt
    - name: Show installed packages
      run: uv pip list
    - name: Test with PyTest
      run: uv run pytest -v -rsx -n 2 -m "compile"

  test_torch_export:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
      uses: astral-sh/setup-uv@v7
      with:
        python-version: "3.14"
        activate-environment: true
    - name: Install dependencies
      run: uv pip install -r requirements/required.txt -r requirements/test.txt
    - name: Show installed packages
      run: uv pip list
    - name: Test with PyTest
      run: uv run pytest -v -rsx -n 2 -m "torch_export"

  test_torch_script:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
      uses: astral-sh/setup-uv@v7
      with:
        python-version: "3.14"
        activate-environment: true
    - name: Install dependencies
      run: uv pip install -r requirements/required.txt -r requirements/test.txt
    - name: Show installed packages
      run: uv pip list
    - name: Test with PyTest
      run: uv run pytest -v -rsx -n 2 -m "torch_script"

  minimum:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
      uses: astral-sh/setup-uv@v7
      with:
        python-version: "3.10"
        activate-environment: true
    - name: Install dependencies
      run: uv pip install -r requirements/minimum.old -r requirements/test.txt
    - name: Show installed packages
      run: uv pip list
    - name: Test with pytest
      run: uv run pytest -v -rsx -n 2 --non-marked-only


================================================
FILE: .gitignore
================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
.idea/
.venv*
examples/images*
examples/annotations*

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
.hypothesis/
.pytest_cache/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py
db.sqlite3

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# Jupyter Notebook
.ipynb_checkpoints

# pyenv
.python-version

# celery beat schedule file
celerybeat-schedule

# SageMath parsed files
*.sage.py

# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
.vscode/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/

# ruff
.ruff_cache/

================================================
FILE: .readthedocs.yaml
================================================
# Read the Docs configuration file for Sphinx projects
# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details

# Required
version: 2

# Set the OS, Python version and other tools you might need
build:
  os: ubuntu-22.04
  tools:
    python: "3.14"

# Build documentation in the "docs/" directory with Sphinx
sphinx:
  configuration: docs/conf.py
  fail_on_warning: true

# Optional but recommended, declare the Python requirements required
# to build your documentation
# See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
python:
  install:
    - requirements: requirements/docs.txt


================================================
FILE: HALLOFFAME.md
================================================
# Hall of Fame

`Segmentation Models` package is widely used in the image segmentation competitions.
Here you can find competitions, names of the winners and links to their solutions.

Please, follow these rules, when adding a solution to the "Hall of Fame":

1. Solution should be high rated (e.g. for Kaggle gold or silver medal) 
2. There should be a description of the solution (post at the forum / code / blog post / paper / pre-print)


## Kaggle

### [Severstal: Steel Defect Detection](https://www.kaggle.com/c/severstal-steel-defect-detection)

- 1st place. 
[Wuxi Jiangsu](https://www.kaggle.com/rguo97), 
[Hongbo Zhu](https://www.kaggle.com/zhuhongbo), 
[Yizhuo Yu](https://www.kaggle.com/paffpaffyu)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114254#latest-675874)]

- 5th place. 
[Guanshuo Xu](https://www.kaggle.com/wowfattie)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/117208#latest-675385)]

- 9th place. 
[Jacek Poplawski](https://www.linkedin.com/in/jacekpoplawski/)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114297#latest-660842)]

- 10th place.
[Alexey Rozhkov](https://www.linkedin.com/in/alexisrozhkov)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114465#latest-659615)]

- 12th place. 
[Pavel Iakubovskii](https://www.linkedin.com/in/pavel-iakubovskii/), 
[Ilya Dobrynin](https://www.linkedin.com/in/ilya-dobrynin-79a89b106/), 
[Denis Kolpakov](https://www.linkedin.com/in/denis-kolpakov-ab3137197/)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114309#latest-661404)]

- 31st place. 
[Insaf Ashrapov](https://www.linkedin.com/in/iashrapov/), 
[Igor Krashenyi](https://www.linkedin.com/in/igor-krashenyi-38b89b98), 
[Pavel Pleskov](https://www.linkedin.com/in/ppleskov), 
[Anton Zakharenkov](https://www.linkedin.com/in/anton-zakharenkov/), 
[Nikolai Popov](https://www.linkedin.com/in/nikolai-popov-b2157370/)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114383#latest-658438)]
[[code](https://github.com/Diyago/Severstal-Steel-Defect-Detection)]

- 55th place. 
[Karl Hornlund](https://www.linkedin.com/in/karl-hornlund/)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/114410#latest-672682)]
[[code](https://github.com/khornlund/severstal-steel-defect-detection)]

- Efficiency round 1st place.
[Stefan Stefanov](https://www.linkedin.com/in/stefan-stefanov-63a77b1)  
[[description](https://www.kaggle.com/c/severstal-steel-defect-detection/discussion/117486#latest-674229)]


### [Understanding Clouds from Satellite Images](https://www.kaggle.com/c/understanding_cloud_organization)

- 2nd place.
[Andrey Kiryasov](https://www.kaggle.com/ekydna)  
[[description](https://www.kaggle.com/c/understanding_cloud_organization/discussion/118255#latest-678189)]

- 4th place.
[Ching-Loong Seow](https://www.linkedin.com/in/clseow/)  
[[description](https://www.kaggle.com/c/understanding_cloud_organization/discussion/118016#latest-677333)]

- 34th place.
[Karl Hornlund](https://www.linkedin.com/in/karl-hornlund/)  
[[description](https://www.kaggle.com/c/understanding_cloud_organization/discussion/118250#latest-678176)]
[[code](https://github.com/khornlund/understanding-cloud-organization)]

- 55th place.
[Pavel Iakubovskii](https://www.linkedin.com/in/pavel-iakubovskii/)  
[[description](https://www.kaggle.com/c/understanding_cloud_organization/discussion/118019#latest-678626)]

## Other platforms

### [MICCAI 2020 TN-SCUI challenge](https://tn-scui2020.grand-challenge.org/Home/)  
- 1st place.
[Mingyu Wang](https://github.com/WAMAWAMA)  
[[description](https://github.com/WAMAWAMA/TNSCUI2020-Seg-Rank1st)]
[[code](https://github.com/WAMAWAMA/TNSCUI2020-Seg-Rank1st)]

### [Open Cities AI Challenge: Segmenting Buildings for Disaster Resilience](https://www.drivendata.org/competitions/60/building-segmentation-disaster-resilience/)
 - 1st place.
[Pavel Iakubovskii](https://www.linkedin.com/in/pavel-iakubovskii/).  
[[code and description](https://github.com/qubvel/open-cities-challenge)]

### [Machine Learning based feature extraction of Electrical Substations from Satellite Data ](https://competitions.codalab.org/competitions/32132#learn_the_details)

- 3rd place.

[Aarsh chaube](https://github.com/Aarsh2001)
[[code](https://github.com/Aarsh2001/ML_Challenge_NRSC)]
[[Pre-Print](https://github.com/Aarsh2001/ML_Challenge_NRSC/blob/main/3rd%20Rank%20Submission.pdf)]

### [NeurIPS2022 Cell Segmentation Challenge](https://neurips22-cellseg.grand-challenge.org/)

- 1st place. [Gihun Lee](https://github.com/Lee-Gihun), [Sangmook Kim](https://github.com/ElvinKim), [Joonkee Kim](https://github.com/joonkeekim)
- [[code](https://github.com/Lee-Gihun/MEDIAR)]
[[Paper](https://arxiv.org/abs/2212.03465)]

### [ICIP2025 Colorectal Cancer Tumor Grade Segmentation in Digital Histopathology Images: From Giga to Mini Challenge](https://sites.google.com/view/cctgs-challenge/home)

- 2nd place. [Ümit Mert Çağlar](linkedin.com/in/mertcaglar), [Alptekin Temizel](https://blog.metu.edu.tr/atemizel/)
- [[code](https://github.com/caglarmert/ICIP2025)]
[[paper](https://doi.org/10.1117/12.3096537)]
[[challange results](https://arxiv.org/abs/2507.04681)]


================================================
FILE: LICENSE
================================================
The MIT License

Copyright (c) 2019, Pavel Iakubovskii

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.


================================================
FILE: Makefile
================================================
.PHONY: test  # Declare the 'test' target as phony to avoid conflicts with files named 'test'

# Variables to store the paths of the python, pip, pytest, and ruff executables
PYTHON := $(shell which python)
PIP := $(shell which pip)
PYTEST := $(shell which pytest)
RUFF := $(shell which ruff)

# Target to create a Python virtual environment
.venv:
	$(PYTHON) -m venv $(shell dirname $(PYTHON))

# Target to install development dependencies in the virtual environment
install_dev: .venv
	$(PIP) install -e ".[test]"

# Target to run tests with pytest, using 2 parallel processes and only non-marked tests
test: .venv
	$(PYTEST) -v -rsx -n 2 tests/ --non-marked-only

# Target to run all tests with pytest, including slow tests, using 2 parallel processes
test_all: .venv
	RUN_SLOW=1 $(PYTEST) -v -rsx -n 2 tests/

# Target to generate a table by running a Python script
table:
	$(PYTHON) misc/generate_table.py

# Target to generate a table for timm by running a Python script
table_timm:
	$(PYTHON) misc/generate_table_timm.py

# Target to fix and format code using ruff
fixup:
	$(RUFF) check --fix
	$(RUFF) format

# Target to run code formatting and tests
all: fixup test


================================================
FILE: README.md
================================================
<div align="center">
 
![logo](https://i.ibb.co/dc1XdhT/Segmentation-Models-V2-Side-1-1.png)  
**Python library with Neural Networks for Image Semantic  
Segmentation based on [PyTorch](https://pytorch.org/).**  

 
[![GitHub Workflow Status (branch)](https://img.shields.io/github/actions/workflow/status/qubvel/segmentation_models.pytorch/tests.yml?branch=main&style=for-the-badge)](https://github.com/qubvel/segmentation_models.pytorch/actions/workflows/tests.yml) 
![Codecov](https://img.shields.io/codecov/c/github/qubvel-org/segmentation_models.pytorch?style=for-the-badge)
[![Read the Docs](https://img.shields.io/readthedocs/smp?style=for-the-badge&logo=readthedocs&logoColor=white)](https://smp.readthedocs.io/en/latest/) 
<br>
[![PyPI](https://img.shields.io/pypi/v/segmentation-models-pytorch?color=red&style=for-the-badge&logo=pypi&logoColor=white)](https://pypi.org/project/segmentation-models-pytorch/) 
[![PyTorch - Version](https://img.shields.io/badge/PYTORCH-1.9+-red?style=for-the-badge&logo=pytorch)](https://pepy.tech/project/segmentation-models-pytorch) 
[![Python - Version](https://img.shields.io/badge/PYTHON-3.9+-red?style=for-the-badge&logo=python&logoColor=white)](https://pepy.tech/project/segmentation-models-pytorch) 
<br>
[![Generic badge](https://img.shields.io/badge/License-MIT-<COLOR>.svg?style=for-the-badge&color=blue)](https://github.com/qubvel/segmentation_models.pytorch/blob/main/LICENSE)
[![PyPI - Downloads](https://img.shields.io/pypi/dm/segmentation-models-pytorch?style=for-the-badge&color=blue)](https://pepy.tech/project/segmentation-models-pytorch) 

</div>

The main features of the library are:

 - Super simple high-level API (just two lines to create a neural network)
 - 12 encoder-decoder model architectures (Unet, Unet++, Segformer, DPT, ...)
 - 800+ **pretrained** convolution- and transform-based encoders, including [timm](https://github.com/huggingface/pytorch-image-models) support
 - Popular metrics and losses for training routines (Dice, Jaccard, Tversky, ...)
 - ONNX export and torch script/trace/compile friendly

### 🤝 Sponsor: withoutBG

<a href="https://github.com/withoutbg/withoutbg" target="_blank">
  <img src="https://withoutbg.com/images/img-and-matte.png" alt="Sponsored by withoutBG"/>
</a>

[withoutBG](https://github.com/withoutbg/withoutbg) is a high-quality background removal tool. They built their open-source image matting and refiner models using `smp.Unet` and are proudly sponsoring this project.


### [📚 Project Documentation 📚](http://smp.readthedocs.io/)

Visit [Read The Docs Project Page](https://smp.readthedocs.io/) or read the following README to know more about Segmentation Models Pytorch (SMP for short) library

### 📋 Table of content
 1. [Quick start](#start)
 2. [Examples](#examples)
 3. [Models and encoders](#models-and-encoders)
 4. [Models API](#api)
    1. [Input channels](#input-channels)
    2. [Auxiliary classification output](#auxiliary-classification-output)
    3. [Depth](#depth)
 5. [Installation](#installation)
 6. [Competitions won with the library](#competitions)
 7. [Contributing](#contributing)
 8. [Citing](#citing)
 9. [License](#license)

## ⏳ Quick start <a name="start"></a>

#### 1. Create your first Segmentation model with SMP

The segmentation model is just a PyTorch `torch.nn.Module`, which can be created as easy as:

```python
import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=3,                      # model output channels (number of classes in your dataset)
)
```
 - see [table](#architectures) with available model architectures
 - see [table](#encoders) with available encoders and their corresponding weights

#### 2. Configure data preprocessing

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give you better results (higher metric score and faster convergence). It is **not necessary** in case you train the whole model, not only the decoder.

```python
from segmentation_models_pytorch.encoders import get_preprocessing_fn

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>

| Name                                      | Link                                                                                           | Colab                                                                                           |
|-------------------------------------------|-----------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|
| **Train** pets binary segmentation on OxfordPets     | [Notebook](https://github.com/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb) |
| **Train** cars binary segmentation on CamVid       | [Notebook](https://github.com/qubvel/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb) |
| **Train** multiclass segmentation on CamVid          | [Notebook](https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/camvid_segmentation_multiclass.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel-org/segmentation_models.pytorch/blob/main/examples/camvid_segmentation_multiclass.ipynb) |
| **Train** clothes binary segmentation by @ternaus   | [Repo](https://github.com/ternaus/cloths_segmentation)                                        |     |
| **Load and inference** pretrained Segformer | [Notebook](https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/segformer_inference_pretrained.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/segformer_inference_pretrained.ipynb) |
| **Load and inference** pretrained DPT | [Notebook](https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/dpt_inference_pretrained.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/dpt_inference_pretrained.ipynb) |
| **Load and inference** pretrained UPerNet | [Notebook](https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/upernet_inference_pretrained.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/upernet_inference_pretrained.ipynb) |
| **Save and load** models locally / to HuggingFace Hub |[Notebook](https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/save_load_model_and_share_with_hf_hub.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/save_load_model_and_share_with_hf_hub.ipynb)
| **Export** trained model to ONNX              | [Notebook](https://github.com/qubvel/segmentation_models.pytorch/blob/main/examples/convert_to_onnx.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/convert_to_onnx.ipynb) |


## 📦 Models and encoders <a name="models-and-encoders"></a>

### Architectures <a name="architectures"></a>
| Architecture | Paper | Documentation | Checkpoints |
|--------------|-------|---------------|------------|
| Unet | [paper](https://arxiv.org/abs/1505.04597) | [docs](https://smp.readthedocs.io/en/latest/models.html#unet) | |
| Unet++ | [paper](https://arxiv.org/pdf/1807.10165.pdf) | [docs](https://smp.readthedocs.io/en/latest/models.html#unetplusplus) | |
| MAnet | [paper](https://ieeexplore.ieee.org/abstract/document/9201310) | [docs](https://smp.readthedocs.io/en/latest/models.html#manet) | |
| Linknet | [paper](https://arxiv.org/abs/1707.03718) | [docs](https://smp.readthedocs.io/en/latest/models.html#linknet) | |
| FPN | [paper](http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf) | [docs](https://smp.readthedocs.io/en/latest/models.html#fpn) | |
| PSPNet | [paper](https://arxiv.org/abs/1612.01105) | [docs](https://smp.readthedocs.io/en/latest/models.html#pspnet) | |
| PAN | [paper](https://arxiv.org/abs/1805.10180) | [docs](https://smp.readthedocs.io/en/latest/models.html#pan) | |
| DeepLabV3 | [paper](https://arxiv.org/abs/1706.05587) | [docs](https://smp.readthedocs.io/en/latest/models.html#deeplabv3) | |
| DeepLabV3+ | [paper](https://arxiv.org/abs/1802.02611) | [docs](https://smp.readthedocs.io/en/latest/models.html#deeplabv3plus) | |
| UPerNet | [paper](https://arxiv.org/abs/1807.10221) | [docs](https://smp.readthedocs.io/en/latest/models.html#upernet) | [checkpoints](https://huggingface.co/collections/smp-hub/upernet-67fadcdbe08418c6ea94f768) |
| Segformer | [paper](https://arxiv.org/abs/2105.15203) | [docs](https://smp.readthedocs.io/en/latest/models.html#segformer) | [checkpoints](https://huggingface.co/collections/smp-hub/segformer-6749eb4923dea2c355f29a1f) |
| DPT | [paper](https://arxiv.org/abs/2103.13413) | [docs](https://smp.readthedocs.io/en/latest/models.html#dpt) | [checkpoints](https://huggingface.co/collections/smp-hub/dpt-67f30487327c0599a0c62d68) |

### Encoders <a name="encoders"></a>

The library provides a wide range of **pretrained** encoders (also known as backbones) for segmentation models. Instead of using features from the final layer of a classification model, we extract **intermediate features** and feed them into the decoder for segmentation tasks.  

All encoders come with **pretrained weights**, which help achieve **faster and more stable convergence** when training segmentation models.  

Given the extensive selection of supported encoders, you can choose the best one for your specific use case, for example:  
- **Lightweight encoders** for low-latency applications or real-time inference on edge devices (mobilenet/mobileone).  
- **High-capacity architectures** for complex tasks involving a large number of segmented classes, providing superior accuracy (convnext/swin/mit).  

By selecting the right encoder, you can balance **efficiency, performance, and model complexity** to suit your project needs.  

All encoders and corresponding pretrained weight are listed in the documentation:
 - [table](https://smp.readthedocs.io/en/latest/encoders.html) with natively ported encoders 
 - [table](https://smp.readthedocs.io/en/latest/encoders_timm.html) with [timm](https://github.com/huggingface/pytorch-image-models) encoders supported

## 🔁 Models API <a name="api"></a>

### Input channels

The input channels parameter allows you to create a model that can process a tensor with an arbitrary number of channels. 
If you use pretrained weights from ImageNet, the weights of the first convolution will be reused:
 - For the 1-channel case, it would be a sum of the weights of the first convolution layer.
 - Otherwise, channels would be populated with weights like `new_weight[:, i] = pretrained_weight[:, i % 3]`, and then scaled with `new_weight * 3 / new_in_channels`.

```python
model = smp.FPN('resnet34', in_channels=1)
mask = model(torch.ones([1, 1, 64, 64]))
```

### Auxiliary classification output

All models support `aux_params` parameters, which is default set to `None`. 
If `aux_params = None` then classification auxiliary output is not created, else
model produce not only `mask`, but also `label` output with shape `NC`.
Classification head consists of GlobalPooling->Dropout(optional)->Linear->Activation(optional) layers, which can be 
configured by `aux_params` as follows:
```python
aux_params=dict(
    pooling='avg',             # one of 'avg', 'max'
    dropout=0.5,               # dropout ratio, default is None
    activation='sigmoid',      # activation function, default is None
    classes=4,                 # define number of output labels
)
model = smp.Unet('resnet34', classes=4, aux_params=aux_params)
mask, label = model(x)
```

### Depth

Depth parameter specify a number of downsampling operations in encoder, so you can make
your model lighter if specify smaller `depth`.
```python
model = smp.Unet('resnet34', encoder_depth=4)
```

## 🛠 Installation <a name="installation"></a>
PyPI version:

```bash
$ pip install segmentation-models-pytorch
````

The latest version from GitHub:

```bash
$ pip install git+https://github.com/qubvel/segmentation_models.pytorch
````

## 🏆 Competitions won with the library <a name="competitions"></a>

`Segmentation Models` package is widely used in image segmentation competitions.
[Here](https://github.com/qubvel/segmentation_models.pytorch/blob/main/HALLOFFAME.md) you can find competitions, names of the winners and links to their solutions.

## 🛠 Projects built with SMP

- [withoutBG](https://github.com/withoutbg/withoutbg): An open-source background removal tool that uses `smp.Unet` for its image matting and refinement models. Check [withoutBG Focus on HuggingFace](https://huggingface.co/withoutbg/focus).


## 🤝 Contributing <a name="contributing"></a>

1. Install SMP in dev mode

```bash
make install_dev  # Create .venv, install SMP in dev mode
```

2. Run tests and code checks

```bash
make test          # Run tests suite with pytest
make fixup         # Ruff for formatting and lint checks
```

3. Update a table (in case you added an encoder)

```bash
make table        # Generates a table with encoders and print to stdout
```

## 📝 Citing <a name="citing"></a>
```
@misc{Iakubovskii:2019,
  Author = {Pavel Iakubovskii},
  Title = {Segmentation Models Pytorch},
  Year = {2019},
  Publisher = {GitHub},
  Journal = {GitHub repository},
  Howpublished = {\url{https://github.com/qubvel/segmentation_models.pytorch}}
}
```

## 🛡️ License <a name="license"></a>
The project is primarily distributed under [MIT License](https://github.com/qubvel/segmentation_models.pytorch/blob/main/LICENSE), while some files are subject to other licenses. Please refer to [LICENSES](licenses/LICENSES.md) and license statements in each file for careful check, especially for commercial use.




================================================
FILE: docs/Makefile
================================================
# Minimal makefile for Sphinx documentation
#

# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS    ?=
SPHINXBUILD   ?= sphinx-build
SOURCEDIR     = .
BUILDDIR      = build

# Put it first so that "make" without argument is like "make help".
help:
	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

.PHONY: help Makefile

# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)


================================================
FILE: docs/conf.py
================================================
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
# sys.path.insert(0, os.path.abspath('.'))

import sys
import datetime
# import sphinx_rtd_theme

sys.path.append("..")

# -- Project information -----------------------------------------------------

project = "Segmentation Models"
copyright = "{}, Pavel Iakubovskii".format(datetime.datetime.now().year)
author = "Pavel Iakubovskii"


def get_version():
    sys.path.append("../segmentation_models_pytorch")
    from __version__ import __version__ as version

    sys.path.pop(-1)
    return version


version = get_version()

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.

extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.coverage",
    "sphinx.ext.intersphinx",
    "sphinx.ext.napoleon",
    "sphinx.ext.viewcode",
    "sphinx.ext.mathjax",
]

# Warn about all references where the target cannot be found.
nitpicky = True
nitpick_ignore = [
    # Undocumented classes
    ("py:class", "torch.FloatTensor"),
    ("py:class", "torch.LongTensor"),
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = []


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#

# html_theme = "sphinx_rtd_theme"
# html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]

# import karma_sphinx_theme
# html_theme = "karma_sphinx_theme"

html_theme = "sphinx_book_theme"

# import catalyst_sphinx_theme
# html_theme = "catalyst_sphinx_theme"
# html_theme_path = [catalyst_sphinx_theme.get_html_theme_path()]

html_logo = "logo.png"

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
# html_static_path = ["_static"]

# -- Extension configuration -------------------------------------------------

autodoc_inherit_docstrings = False
napoleon_google_docstring = True
napoleon_include_init_with_doc = True
napoleon_numpy_docstring = False

autodoc_mock_imports = [
    "torch",
    "tqdm",
    "numpy",
    "timm",
    "cv2",
    "PIL",
    "torchvision",
    "segmentation_models_pytorch.encoders",
    "segmentation_models_pytorch.utils",
    # 'segmentation_models_pytorch.base',
]

autoclass_content = "both"
autodoc_typehints = "description"

# sphinx.ext.intersphinx
intersphinx_mapping = {
    "python": ("https://docs.python.org/3", None),
    "torch": ("https://docs.pytorch.org/docs/stable/", None),
}

# --- Work around to make autoclass signatures not (*args, **kwargs) ----------


class FakeSignature:
    def __getattribute__(self, *args):
        raise ValueError


def f(app, obj, bound_method):
    if "__new__" in obj.__name__:
        obj.__signature__ = FakeSignature()


def setup(app):
    app.connect("autodoc-before-process-signature", f)


# Custom configuration --------------------------------------------------------

autodoc_member_order = "bysource"


================================================
FILE: docs/encoders.rst
================================================
🔍 Available Encoders
=====================

**Segmentation Models PyTorch** provides support for a wide range of encoders. 
This flexibility allows you to use these encoders with any model in the library by 
specifying the encoder name in the ``encoder_name`` parameter during model initialization.

Here’s a quick example of using a ResNet34 encoder with the ``Unet`` model:

.. code-block:: python

    from segmentation_models_pytorch import Unet

    # Initialize Unet with ResNet34 encoder pre-trained on ImageNet
    model = Unet(encoder_name="resnet34", encoder_weights="imagenet")


The following encoder families are supported by the library, enabling you to choose the one that best fits your use case:

- **Mix Vision Transformer (mit)**
- **MobileOne**
- **MobileNet**
- **EfficientNet**
- **ResNet**
- **ResNeXt**
- **SENet**
- **DPN**
- **VGG**
- **DenseNet**
- **Xception**
- **Inception**

Choosing the Right Encoder
--------------------------

1. **Small Models for Edge Devices**  
   Consider encoders like **MobileNet** or **MobileOne**, which have a smaller parameter count and are optimized for lightweight deployment.

2. **High Performance**  
   If you require state-of-the-art accuracy **Mix Vision Transformer (mit)**, **EfficientNet** families offer excellent balance between performance and computational efficiency.

For each encoder, the table below provides detailed information:

1. **Pretrained Weights**  
   Specifies the available pretrained weights (e.g., ``imagenet``, ``imagenet21k``).

2. **Params, M**: 
   The total number of parameters in the encoder, measured in millions. This metric helps you assess the model's size and computational requirements.

3. **Script**:  
   Indicates whether the encoder can be scripted with ``torch.jit.script``.

4. **Compile**:   
   Indicates whether the encoder is compatible with ``torch.compile(model, fullgraph=True, dynamic=True, backend="eager")``. 
   You may still get some issues with another backends, such as ``inductor``, depending on the torch/cuda/... dependencies version,
   but most of the time it will work.

5. **Export**:   
   Indicates whether the encoder can be exported using ``torch.export.export``, making it suitable for deployment in different environments (e.g., ONNX).


============================ ==================================== =========== ======== ========= ======== 
Encoder                      Pretrained weights                   Params, M   Script   Compile   Export  
============================ ==================================== =========== ======== ========= ======== 
resnet18                     imagenet / ssl / swsl                11M         ✅        ✅         ✅       
resnet34                     imagenet                             21M         ✅        ✅         ✅       
resnet50                     imagenet / ssl / swsl                23M         ✅        ✅         ✅       
resnet101                    imagenet                             42M         ✅        ✅         ✅       
resnet152                    imagenet                             58M         ✅        ✅         ✅       
resnext50_32x4d              imagenet / ssl / swsl                22M         ✅        ✅         ✅       
resnext101_32x4d             ssl / swsl                           42M         ✅        ✅         ✅       
resnext101_32x8d             imagenet / instagram / ssl / swsl    86M         ✅        ✅         ✅       
resnext101_32x16d            instagram / ssl / swsl               191M        ✅        ✅         ✅       
resnext101_32x32d            instagram                            466M        ✅        ✅         ✅       
resnext101_32x48d            instagram                            826M        ✅        ✅         ✅       
dpn68                        imagenet                             11M         ❌        ✅         ✅       
dpn68b                       imagenet+5k                          11M         ❌        ✅         ✅       
dpn92                        imagenet+5k                          34M         ❌        ✅         ✅       
dpn98                        imagenet                             58M         ❌        ✅         ✅       
dpn107                       imagenet+5k                          84M         ❌        ✅         ✅       
dpn131                       imagenet                             76M         ❌        ✅         ✅       
vgg11                        imagenet                             9M          ✅        ✅         ✅       
vgg11_bn                     imagenet                             9M          ✅        ✅         ✅       
vgg13                        imagenet                             9M          ✅        ✅         ✅       
vgg13_bn                     imagenet                             9M          ✅        ✅         ✅       
vgg16                        imagenet                             14M         ✅        ✅         ✅       
vgg16_bn                     imagenet                             14M         ✅        ✅         ✅       
vgg19                        imagenet                             20M         ✅        ✅         ✅       
vgg19_bn                     imagenet                             20M         ✅        ✅         ✅       
senet154                     imagenet                             113M        ✅        ✅         ✅       
se_resnet50                  imagenet                             26M         ✅        ✅         ✅       
se_resnet101                 imagenet                             47M         ✅        ✅         ✅       
se_resnet152                 imagenet                             64M         ✅        ✅         ✅       
se_resnext50_32x4d           imagenet                             25M         ✅        ✅         ✅       
se_resnext101_32x4d          imagenet                             46M         ✅        ✅         ✅       
densenet121                  imagenet                             6M          ✅        ✅         ✅       
densenet169                  imagenet                             12M         ✅        ✅         ✅       
densenet201                  imagenet                             18M         ✅        ✅         ✅       
densenet161                  imagenet                             26M         ✅        ✅         ✅       
inceptionresnetv2            imagenet / imagenet+background       54M         ✅        ✅         ✅       
inceptionv4                  imagenet / imagenet+background       41M         ✅        ✅         ✅       
efficientnet-b0              imagenet / advprop                   4M          ✅        ✅         ✅       
efficientnet-b1              imagenet / advprop                   6M          ✅        ✅         ✅       
efficientnet-b2              imagenet / advprop                   7M          ✅        ✅         ✅       
efficientnet-b3              imagenet / advprop                   10M         ✅        ✅         ✅       
efficientnet-b4              imagenet / advprop                   17M         ✅        ✅         ✅       
efficientnet-b5              imagenet / advprop                   28M         ✅        ✅         ✅       
efficientnet-b6              imagenet / advprop                   40M         ✅        ✅         ✅       
efficientnet-b7              imagenet / advprop                   63M         ✅        ✅         ✅       
mobilenet_v2                 imagenet                             2M          ✅        ✅         ✅       
xception                     imagenet                             20M         ✅        ✅         ✅       
timm-efficientnet-b0         imagenet / advprop / noisy-student   4M          ✅        ✅         ✅       
timm-efficientnet-b1         imagenet / advprop / noisy-student   6M          ✅        ✅         ✅       
timm-efficientnet-b2         imagenet / advprop / noisy-student   7M          ✅        ✅         ✅       
timm-efficientnet-b3         imagenet / advprop / noisy-student   10M         ✅        ✅         ✅       
timm-efficientnet-b4         imagenet / advprop / noisy-student   17M         ✅        ✅         ✅       
timm-efficientnet-b5         imagenet / advprop / noisy-student   28M         ✅        ✅         ✅       
timm-efficientnet-b6         imagenet / advprop / noisy-student   40M         ✅        ✅         ✅       
timm-efficientnet-b7         imagenet / advprop / noisy-student   63M         ✅        ✅         ✅       
timm-efficientnet-b8         imagenet / advprop                   84M         ✅        ✅         ✅       
timm-efficientnet-l2         noisy-student / noisy-student-475    474M        ✅        ✅         ✅       
timm-tf_efficientnet_lite0   imagenet                             3M          ✅        ✅         ✅       
timm-tf_efficientnet_lite1   imagenet                             4M          ✅        ✅         ✅       
timm-tf_efficientnet_lite2   imagenet                             4M          ✅        ✅         ✅       
timm-tf_efficientnet_lite3   imagenet                             6M          ✅        ✅         ✅       
timm-tf_efficientnet_lite4   imagenet                             11M         ✅        ✅         ✅       
timm-skresnet18              imagenet                             11M         ✅        ✅         ✅       
timm-skresnet34              imagenet                             21M         ✅        ✅         ✅       
timm-skresnext50_32x4d       imagenet                             23M         ✅        ✅         ✅       
mit_b0                       imagenet                             3M          ✅        ✅         ✅       
mit_b1                       imagenet                             13M         ✅        ✅         ✅       
mit_b2                       imagenet                             24M         ✅        ✅         ✅       
mit_b3                       imagenet                             44M         ✅        ✅         ✅       
mit_b4                       imagenet                             60M         ✅        ✅         ✅       
mit_b5                       imagenet                             81M         ✅        ✅         ✅       
mobileone_s0                 imagenet                             4M          ✅        ✅         ✅       
mobileone_s1                 imagenet                             3M          ✅        ✅         ✅       
mobileone_s2                 imagenet                             5M          ✅        ✅         ✅       
mobileone_s3                 imagenet                             8M          ✅        ✅         ✅       
mobileone_s4                 imagenet                             12M         ✅        ✅         ✅       
============================ ==================================== =========== ======== ========= ======== 


================================================
FILE: docs/encoders_dpt.rst
================================================
:orphan:

.. _dpt-encoders:

DPT Encoders
============

This is a list of Vision Transformer encoders that are compatible with the DPT architecture.
While other Vision Transformer encoders from timm may also be compatible, the ones listed below are tested to work properly.

.. list-table:: Encoder Name
   :widths: 100
   :header-rows: 0

   * - tu-fastvit_ma36.apple_dist_in1k
   * - tu-fastvit_ma36.apple_in1k
   * - tu-fastvit_mci0.apple_mclip
   * - tu-fastvit_mci1.apple_mclip
   * - tu-fastvit_mci2.apple_mclip
   * - tu-fastvit_s12.apple_dist_in1k
   * - tu-fastvit_s12.apple_in1k
   * - tu-fastvit_sa12.apple_dist_in1k
   * - tu-fastvit_sa12.apple_in1k
   * - tu-fastvit_sa24.apple_dist_in1k
   * - tu-fastvit_sa24.apple_in1k
   * - tu-fastvit_sa36.apple_dist_in1k
   * - tu-fastvit_sa36.apple_in1k
   * - tu-fastvit_t8.apple_dist_in1k
   * - tu-fastvit_t8.apple_in1k
   * - tu-fastvit_t12.apple_dist_in1k
   * - tu-fastvit_t12.apple_in1k
   * - tu-flexivit_base.300ep_in1k
   * - tu-flexivit_base.300ep_in21k
   * - tu-flexivit_base.600ep_in1k
   * - tu-flexivit_base.1000ep_in21k
   * - tu-flexivit_base.1200ep_in1k
   * - tu-flexivit_base.patch16_in21k
   * - tu-flexivit_base.patch30_in21k
   * - tu-flexivit_large.300ep_in1k
   * - tu-flexivit_large.600ep_in1k
   * - tu-flexivit_large.1200ep_in1k
   * - tu-flexivit_small.300ep_in1k
   * - tu-flexivit_small.600ep_in1k
   * - tu-flexivit_small.1200ep_in1k
   * - tu-maxvit_base_tf_224.in1k
   * - tu-maxvit_base_tf_224.in21k
   * - tu-maxvit_base_tf_384.in1k
   * - tu-maxvit_base_tf_384.in21k_ft_in1k
   * - tu-maxvit_base_tf_512.in1k
   * - tu-maxvit_base_tf_512.in21k_ft_in1k
   * - tu-maxvit_large_tf_224.in1k
   * - tu-maxvit_large_tf_224.in21k
   * - tu-maxvit_large_tf_384.in1k
   * - tu-maxvit_large_tf_384.in21k_ft_in1k
   * - tu-maxvit_large_tf_512.in1k
   * - tu-maxvit_large_tf_512.in21k_ft_in1k
   * - tu-maxvit_nano_rw_256.sw_in1k
   * - tu-maxvit_rmlp_base_rw_224.sw_in12k
   * - tu-maxvit_rmlp_base_rw_224.sw_in12k_ft_in1k
   * - tu-maxvit_rmlp_base_rw_384.sw_in12k_ft_in1k
   * - tu-maxvit_rmlp_nano_rw_256.sw_in1k
   * - tu-maxvit_rmlp_pico_rw_256.sw_in1k
   * - tu-maxvit_rmlp_small_rw_224.sw_in1k
   * - tu-maxvit_rmlp_tiny_rw_256.sw_in1k
   * - tu-maxvit_small_tf_224.in1k
   * - tu-maxvit_small_tf_384.in1k
   * - tu-maxvit_small_tf_512.in1k
   * - tu-maxvit_tiny_rw_224.sw_in1k
   * - tu-maxvit_tiny_tf_224.in1k
   * - tu-maxvit_tiny_tf_384.in1k
   * - tu-maxvit_tiny_tf_512.in1k
   * - tu-maxvit_xlarge_tf_224.in21k
   * - tu-maxvit_xlarge_tf_384.in21k_ft_in1k
   * - tu-maxvit_xlarge_tf_512.in21k_ft_in1k
   * - tu-maxxvit_rmlp_nano_rw_256.sw_in1k
   * - tu-maxxvit_rmlp_small_rw_256.sw_in1k
   * - tu-maxxvitv2_nano_rw_256.sw_in1k
   * - tu-maxxvitv2_rmlp_base_rw_224.sw_in12k
   * - tu-maxxvitv2_rmlp_base_rw_224.sw_in12k_ft_in1k
   * - tu-maxxvitv2_rmlp_base_rw_384.sw_in12k_ft_in1k
   * - tu-mobilevit_s.cvnets_in1k
   * - tu-mobilevit_xs.cvnets_in1k
   * - tu-mobilevit_xxs.cvnets_in1k
   * - tu-mobilevitv2_050.cvnets_in1k
   * - tu-mobilevitv2_075.cvnets_in1k
   * - tu-mobilevitv2_100.cvnets_in1k
   * - tu-mobilevitv2_125.cvnets_in1k
   * - tu-mobilevitv2_150.cvnets_in1k
   * - tu-mobilevitv2_150.cvnets_in22k_ft_in1k
   * - tu-mobilevitv2_150.cvnets_in22k_ft_in1k_384
   * - tu-mobilevitv2_175.cvnets_in1k
   * - tu-mobilevitv2_175.cvnets_in22k_ft_in1k
   * - tu-mobilevitv2_175.cvnets_in22k_ft_in1k_384
   * - tu-mobilevitv2_200.cvnets_in1k
   * - tu-mobilevitv2_200.cvnets_in22k_ft_in1k
   * - tu-mobilevitv2_200.cvnets_in22k_ft_in1k_384
   * - tu-mvitv2_base.fb_in1k
   * - tu-mvitv2_base_cls.fb_inw21k
   * - tu-mvitv2_huge_cls.fb_inw21k
   * - tu-mvitv2_large.fb_in1k
   * - tu-mvitv2_large_cls.fb_inw21k
   * - tu-mvitv2_small.fb_in1k
   * - tu-mvitv2_tiny.fb_in1k
   * - tu-samvit_base_patch16.sa1b
   * - tu-samvit_huge_patch16.sa1b
   * - tu-samvit_large_patch16.sa1b
   * - tu-test_vit2.r160_in1k
   * - tu-test_vit3.r160_in1k
   * - tu-test_vit.r160_in1k
   * - tu-vit_base_mci_224.apple_mclip
   * - tu-vit_base_mci_224.apple_mclip_lt
   * - tu-vit_base_patch8_224.augreg2_in21k_ft_in1k
   * - tu-vit_base_patch8_224.augreg_in21k
   * - tu-vit_base_patch8_224.augreg_in21k_ft_in1k
   * - tu-vit_base_patch8_224.dino
   * - tu-vit_base_patch16_224.augreg2_in21k_ft_in1k
   * - tu-vit_base_patch16_224.augreg_in1k
   * - tu-vit_base_patch16_224.augreg_in21k
   * - tu-vit_base_patch16_224.augreg_in21k_ft_in1k
   * - tu-vit_base_patch16_224.dino
   * - tu-vit_base_patch16_224.mae
   * - tu-vit_base_patch16_224.orig_in21k
   * - tu-vit_base_patch16_224.orig_in21k_ft_in1k
   * - tu-vit_base_patch16_224.sam_in1k
   * - tu-vit_base_patch16_224_miil.in21k
   * - tu-vit_base_patch16_224_miil.in21k_ft_in1k
   * - tu-vit_base_patch16_384.augreg_in1k
   * - tu-vit_base_patch16_384.augreg_in21k_ft_in1k
   * - tu-vit_base_patch16_384.orig_in21k_ft_in1k
   * - tu-vit_base_patch16_clip_224.datacompxl
   * - tu-vit_base_patch16_clip_224.dfn2b
   * - tu-vit_base_patch16_clip_224.laion2b
   * - tu-vit_base_patch16_clip_224.laion2b_ft_in1k
   * - tu-vit_base_patch16_clip_224.laion2b_ft_in12k
   * - tu-vit_base_patch16_clip_224.laion2b_ft_in12k_in1k
   * - tu-vit_base_patch16_clip_224.laion400m_e32
   * - tu-vit_base_patch16_clip_224.metaclip_2pt5b
   * - tu-vit_base_patch16_clip_224.metaclip_400m
   * - tu-vit_base_patch16_clip_224.openai
   * - tu-vit_base_patch16_clip_224.openai_ft_in1k
   * - tu-vit_base_patch16_clip_224.openai_ft_in12k
   * - tu-vit_base_patch16_clip_224.openai_ft_in12k_in1k
   * - tu-vit_base_patch16_clip_384.laion2b_ft_in1k
   * - tu-vit_base_patch16_clip_384.laion2b_ft_in12k_in1k
   * - tu-vit_base_patch16_clip_384.openai_ft_in1k
   * - tu-vit_base_patch16_clip_384.openai_ft_in12k_in1k
   * - tu-vit_base_patch16_clip_quickgelu_224.metaclip_2pt5b
   * - tu-vit_base_patch16_clip_quickgelu_224.metaclip_400m
   * - tu-vit_base_patch16_clip_quickgelu_224.openai
   * - tu-vit_base_patch16_plus_clip_240.laion400m_e32
   * - tu-vit_base_patch16_rope_reg1_gap_256.sbb_in1k
   * - tu-vit_base_patch16_rpn_224.sw_in1k
   * - tu-vit_base_patch16_siglip_224.v2_webli
   * - tu-vit_base_patch16_siglip_224.webli
   * - tu-vit_base_patch16_siglip_256.v2_webli
   * - tu-vit_base_patch16_siglip_256.webli
   * - tu-vit_base_patch16_siglip_256.webli_i18n
   * - tu-vit_base_patch16_siglip_384.v2_webli
   * - tu-vit_base_patch16_siglip_384.webli
   * - tu-vit_base_patch16_siglip_512.v2_webli
   * - tu-vit_base_patch16_siglip_512.webli
   * - tu-vit_base_patch16_siglip_gap_224.v2_webli
   * - tu-vit_base_patch16_siglip_gap_224.webli
   * - tu-vit_base_patch16_siglip_gap_256.v2_webli
   * - tu-vit_base_patch16_siglip_gap_256.webli
   * - tu-vit_base_patch16_siglip_gap_256.webli_i18n
   * - tu-vit_base_patch16_siglip_gap_384.v2_webli
   * - tu-vit_base_patch16_siglip_gap_384.webli
   * - tu-vit_base_patch16_siglip_gap_512.v2_webli
   * - tu-vit_base_patch16_siglip_gap_512.webli
   * - tu-vit_base_patch32_224.augreg_in1k
   * - tu-vit_base_patch32_224.augreg_in21k
   * - tu-vit_base_patch32_224.augreg_in21k_ft_in1k
   * - tu-vit_base_patch32_224.orig_in21k
   * - tu-vit_base_patch32_224.sam_in1k
   * - tu-vit_base_patch32_384.augreg_in1k
   * - tu-vit_base_patch32_384.augreg_in21k_ft_in1k
   * - tu-vit_base_patch32_clip_224.datacompxl
   * - tu-vit_base_patch32_clip_224.laion2b
   * - tu-vit_base_patch32_clip_224.laion2b_ft_in1k
   * - tu-vit_base_patch32_clip_224.laion2b_ft_in12k_in1k
   * - tu-vit_base_patch32_clip_224.laion400m_e32
   * - tu-vit_base_patch32_clip_224.metaclip_2pt5b
   * - tu-vit_base_patch32_clip_224.metaclip_400m
   * - tu-vit_base_patch32_clip_224.openai
   * - tu-vit_base_patch32_clip_224.openai_ft_in1k
   * - tu-vit_base_patch32_clip_256.datacompxl
   * - tu-vit_base_patch32_clip_384.laion2b_ft_in12k_in1k
   * - tu-vit_base_patch32_clip_384.openai_ft_in12k_in1k
   * - tu-vit_base_patch32_clip_448.laion2b_ft_in12k_in1k
   * - tu-vit_base_patch32_clip_quickgelu_224.laion400m_e32
   * - tu-vit_base_patch32_clip_quickgelu_224.metaclip_2pt5b
   * - tu-vit_base_patch32_clip_quickgelu_224.metaclip_400m
   * - tu-vit_base_patch32_clip_quickgelu_224.openai
   * - tu-vit_base_patch32_siglip_256.v2_webli
   * - tu-vit_base_patch32_siglip_gap_256.v2_webli
   * - tu-vit_base_r50_s16_224.orig_in21k
   * - tu-vit_base_r50_s16_384.orig_in21k_ft_in1k
   * - tu-vit_betwixt_patch16_reg1_gap_256.sbb_in1k
   * - tu-vit_betwixt_patch16_reg4_gap_256.sbb2_e200_in12k
   * - tu-vit_betwixt_patch16_reg4_gap_256.sbb2_e200_in12k_ft_in1k
   * - tu-vit_betwixt_patch16_reg4_gap_256.sbb_in1k
   * - tu-vit_betwixt_patch16_reg4_gap_256.sbb_in12k
   * - tu-vit_betwixt_patch16_reg4_gap_256.sbb_in12k_ft_in1k
   * - tu-vit_betwixt_patch16_reg4_gap_384.sbb2_e200_in12k_ft_in1k
   * - tu-vit_betwixt_patch16_rope_reg4_gap_256.sbb_in1k
   * - tu-vit_betwixt_patch32_clip_224.tinyclip_laion400m
   * - tu-vit_giant_patch16_gap_224.in22k_ijepa
   * - tu-vit_giantopt_patch16_siglip_256.v2_webli
   * - tu-vit_giantopt_patch16_siglip_384.v2_webli
   * - tu-vit_giantopt_patch16_siglip_gap_256.v2_webli
   * - tu-vit_giantopt_patch16_siglip_gap_384.v2_webli
   * - tu-vit_huge_patch16_gap_448.in1k_ijepa
   * - tu-vit_large_patch16_224.augreg_in21k
   * - tu-vit_large_patch16_224.augreg_in21k_ft_in1k
   * - tu-vit_large_patch16_224.mae
   * - tu-vit_large_patch16_224.orig_in21k
   * - tu-vit_large_patch16_384.augreg_in21k_ft_in1k
   * - tu-vit_large_patch16_siglip_256.v2_webli
   * - tu-vit_large_patch16_siglip_256.webli
   * - tu-vit_large_patch16_siglip_384.v2_webli
   * - tu-vit_large_patch16_siglip_384.webli
   * - tu-vit_large_patch16_siglip_512.v2_webli
   * - tu-vit_large_patch16_siglip_gap_256.v2_webli
   * - tu-vit_large_patch16_siglip_gap_256.webli
   * - tu-vit_large_patch16_siglip_gap_384.v2_webli
   * - tu-vit_large_patch16_siglip_gap_384.webli
   * - tu-vit_large_patch16_siglip_gap_512.v2_webli
   * - tu-vit_large_patch32_224.orig_in21k
   * - tu-vit_large_patch32_384.orig_in21k_ft_in1k
   * - tu-vit_large_r50_s32_224.augreg_in21k
   * - tu-vit_large_r50_s32_224.augreg_in21k_ft_in1k
   * - tu-vit_large_r50_s32_384.augreg_in21k_ft_in1k
   * - tu-vit_little_patch16_reg1_gap_256.sbb_in12k
   * - tu-vit_little_patch16_reg1_gap_256.sbb_in12k_ft_in1k
   * - tu-vit_little_patch16_reg4_gap_256.sbb_in1k
   * - tu-vit_medium_patch16_clip_224.tinyclip_yfcc15m
   * - tu-vit_medium_patch16_gap_240.sw_in12k
   * - tu-vit_medium_patch16_gap_256.sw_in12k_ft_in1k
   * - tu-vit_medium_patch16_gap_384.sw_in12k_ft_in1k
   * - tu-vit_medium_patch16_reg1_gap_256.sbb_in1k
   * - tu-vit_medium_patch16_reg4_gap_256.sbb_in1k
   * - tu-vit_medium_patch16_reg4_gap_256.sbb_in12k
   * - tu-vit_medium_patch16_reg4_gap_256.sbb_in12k_ft_in1k
   * - tu-vit_medium_patch16_rope_reg1_gap_256.sbb_in1k
   * - tu-vit_medium_patch32_clip_224.tinyclip_laion400m
   * - tu-vit_mediumd_patch16_reg4_gap_256.sbb2_e200_in12k
   * - tu-vit_mediumd_patch16_reg4_gap_256.sbb2_e200_in12k_ft_in1k
   * - tu-vit_mediumd_patch16_reg4_gap_256.sbb_in12k
   * - tu-vit_mediumd_patch16_reg4_gap_256.sbb_in12k_ft_in1k
   * - tu-vit_mediumd_patch16_reg4_gap_384.sbb2_e200_in12k_ft_in1k
   * - tu-vit_mediumd_patch16_rope_reg1_gap_256.sbb_in1k
   * - tu-vit_pwee_patch16_reg1_gap_256.sbb_in1k
   * - tu-vit_relpos_base_patch16_224.sw_in1k
   * - tu-vit_relpos_base_patch16_clsgap_224.sw_in1k
   * - tu-vit_relpos_base_patch32_plus_rpn_256.sw_in1k
   * - tu-vit_relpos_medium_patch16_224.sw_in1k
   * - tu-vit_relpos_medium_patch16_cls_224.sw_in1k
   * - tu-vit_relpos_medium_patch16_rpn_224.sw_in1k
   * - tu-vit_relpos_small_patch16_224.sw_in1k
   * - tu-vit_small_patch8_224.dino
   * - tu-vit_small_patch16_224.augreg_in1k
   * - tu-vit_small_patch16_224.augreg_in21k
   * - tu-vit_small_patch16_224.augreg_in21k_ft_in1k
   * - tu-vit_small_patch16_224.dino
   * - tu-vit_small_patch16_384.augreg_in1k
   * - tu-vit_small_patch16_384.augreg_in21k_ft_in1k
   * - tu-vit_small_patch32_224.augreg_in21k
   * - tu-vit_small_patch32_224.augreg_in21k_ft_in1k
   * - tu-vit_small_patch32_384.augreg_in21k_ft_in1k
   * - tu-vit_small_r26_s32_224.augreg_in21k
   * - tu-vit_small_r26_s32_224.augreg_in21k_ft_in1k
   * - tu-vit_small_r26_s32_384.augreg_in21k_ft_in1k
   * - tu-vit_so150m2_patch16_reg1_gap_256.sbb_e200_in12k
   * - tu-vit_so150m2_patch16_reg1_gap_256.sbb_e200_in12k_ft_in1k
   * - tu-vit_so150m2_patch16_reg1_gap_384.sbb_e200_in12k_ft_in1k
   * - tu-vit_so150m2_patch16_reg1_gap_448.sbb_e200_in12k_ft_in1k
   * - tu-vit_so150m_patch16_reg4_gap_256.sbb_e250_in12k
   * - tu-vit_so150m_patch16_reg4_gap_256.sbb_e250_in12k_ft_in1k
   * - tu-vit_so150m_patch16_reg4_gap_384.sbb_e250_in12k_ft_in1k
   * - tu-vit_so400m_patch16_siglip_256.v2_webli
   * - tu-vit_so400m_patch16_siglip_256.webli_i18n
   * - tu-vit_so400m_patch16_siglip_384.v2_webli
   * - tu-vit_so400m_patch16_siglip_512.v2_webli
   * - tu-vit_so400m_patch16_siglip_gap_256.v2_webli
   * - tu-vit_so400m_patch16_siglip_gap_256.webli_i18n
   * - tu-vit_so400m_patch16_siglip_gap_384.v2_webli
   * - tu-vit_so400m_patch16_siglip_gap_512.v2_webli
   * - tu-vit_srelpos_medium_patch16_224.sw_in1k
   * - tu-vit_srelpos_small_patch16_224.sw_in1k
   * - tu-vit_tiny_patch16_224.augreg_in21k
   * - tu-vit_tiny_patch16_224.augreg_in21k_ft_in1k
   * - tu-vit_tiny_patch16_384.augreg_in21k_ft_in1k
   * - tu-vit_tiny_r_s16_p8_224.augreg_in21k
   * - tu-vit_tiny_r_s16_p8_224.augreg_in21k_ft_in1k
   * - tu-vit_tiny_r_s16_p8_384.augreg_in21k_ft_in1k
   * - tu-vit_wee_patch16_reg1_gap_256.sbb_in1k
   * - tu-vit_xsmall_patch16_clip_224.tinyclip_yfcc15m
   * - tu-vitamin_base_224.datacomp1b_clip
   * - tu-vitamin_base_224.datacomp1b_clip_ltt
   * - tu-vitamin_large2_224.datacomp1b_clip
   * - tu-vitamin_large2_256.datacomp1b_clip
   * - tu-vitamin_large2_336.datacomp1b_clip
   * - tu-vitamin_large2_384.datacomp1b_clip
   * - tu-vitamin_large_224.datacomp1b_clip
   * - tu-vitamin_large_256.datacomp1b_clip
   * - tu-vitamin_large_336.datacomp1b_clip
   * - tu-vitamin_large_384.datacomp1b_clip
   * - tu-vitamin_small_224.datacomp1b_clip
   * - tu-vitamin_small_224.datacomp1b_clip_ltt
   * - tu-vitamin_xlarge_256.datacomp1b_clip
   * - tu-vitamin_xlarge_336.datacomp1b_clip
   * - tu-vitamin_xlarge_384.datacomp1b_clip
   * - tu-hiera_small_abswin_256.sbb2_e200_in12k
   * - tu-hiera_small_abswin_256.sbb2_e200_in12k_ft_in1k
   * - tu-hiera_small_abswin_256.sbb2_pd_e200_in12k
   * - tu-hiera_small_abswin_256.sbb2_pd_e200_in12k_ft_in1k
   * - tu-swin_base_patch4_window7_224.ms_in1k
   * - tu-swin_base_patch4_window7_224.ms_in22k
   * - tu-swin_base_patch4_window7_224.ms_in22k_ft_in1k
   * - tu-swin_base_patch4_window12_384.ms_in1k
   * - tu-swin_base_patch4_window12_384.ms_in22k
   * - tu-swin_base_patch4_window12_384.ms_in22k_ft_in1k
   * - tu-swin_large_patch4_window7_224.ms_in22k
   * - tu-swin_large_patch4_window7_224.ms_in22k_ft_in1k
   * - tu-swin_large_patch4_window12_384.ms_in22k
   * - tu-swin_large_patch4_window12_384.ms_in22k_ft_in1k
   * - tu-swin_s3_base_224.ms_in1k
   * - tu-swin_s3_small_224.ms_in1k
   * - tu-swin_s3_tiny_224.ms_in1k
   * - tu-swin_small_patch4_window7_224.ms_in1k
   * - tu-swin_small_patch4_window7_224.ms_in22k
   * - tu-swin_small_patch4_window7_224.ms_in22k_ft_in1k
   * - tu-swin_tiny_patch4_window7_224.ms_in1k
   * - tu-swin_tiny_patch4_window7_224.ms_in22k
   * - tu-swin_tiny_patch4_window7_224.ms_in22k_ft_in1k
   * - tu-swinv2_base_window8_256.ms_in1k
   * - tu-swinv2_base_window12_192.ms_in22k
   * - tu-swinv2_base_window12to16_192to256.ms_in22k_ft_in1k
   * - tu-swinv2_base_window12to24_192to384.ms_in22k_ft_in1k
   * - tu-swinv2_base_window16_256.ms_in1k
   * - tu-swinv2_cr_small_224.sw_in1k
   * - tu-swinv2_cr_small_ns_224.sw_in1k
   * - tu-swinv2_cr_tiny_ns_224.sw_in1k
   * - tu-swinv2_large_window12_192.ms_in22k
   * - tu-swinv2_large_window12to16_192to256.ms_in22k_ft_in1k
   * - tu-swinv2_large_window12to24_192to384.ms_in22k_ft_in1k
   * - tu-swinv2_small_window8_256.ms_in1k
   * - tu-swinv2_small_window16_256.ms_in1k
   * - tu-swinv2_tiny_window8_256.ms_in1k
   * - tu-swinv2_tiny_window16_256.ms_in1k
   * - tu-efficientformer_l1.snap_dist_in1k
   * - tu-efficientformer_l3.snap_dist_in1k
   * - tu-efficientformer_l7.snap_dist_in1k
   * - tu-beit_base_patch16_224.in22k_ft_in22k
   * - tu-beit_base_patch16_224.in22k_ft_in22k_in1k
   * - tu-beit_base_patch16_384.in22k_ft_in22k_in1k
   * - tu-beit_large_patch16_224.in22k_ft_in22k
   * - tu-beit_large_patch16_224.in22k_ft_in22k_in1k
   * - tu-beit_large_patch16_384.in22k_ft_in22k_in1k
   * - tu-beit_large_patch16_512.in22k_ft_in22k_in1k
   * - tu-beitv2_base_patch16_224.in1k_ft_in1k
   * - tu-beitv2_base_patch16_224.in1k_ft_in22k
   * - tu-beitv2_base_patch16_224.in1k_ft_in22k_in1k
   * - tu-beitv2_large_patch16_224.in1k_ft_in1k
   * - tu-beitv2_large_patch16_224.in1k_ft_in22k
   * - tu-beitv2_large_patch16_224.in1k_ft_in22k_in1k
   * - tu-cait_m36_384.fb_dist_in1k
   * - tu-cait_m48_448.fb_dist_in1k
   * - tu-cait_s24_224.fb_dist_in1k
   * - tu-cait_s24_384.fb_dist_in1k
   * - tu-cait_s36_384.fb_dist_in1k
   * - tu-cait_xs24_384.fb_dist_in1k
   * - tu-cait_xxs24_224.fb_dist_in1k
   * - tu-cait_xxs24_384.fb_dist_in1k
   * - tu-cait_xxs36_224.fb_dist_in1k
   * - tu-cait_xxs36_384.fb_dist_in1k
   * - tu-coatnet_0_rw_224.sw_in1k
   * - tu-coatnet_1_rw_224.sw_in1k
   * - tu-coatnet_2_rw_224.sw_in12k
   * - tu-coatnet_2_rw_224.sw_in12k_ft_in1k
   * - tu-coatnet_3_rw_224.sw_in12k
   * - tu-coatnet_bn_0_rw_224.sw_in1k
   * - tu-coatnet_nano_rw_224.sw_in1k
   * - tu-coatnet_rmlp_1_rw2_224.sw_in12k
   * - tu-coatnet_rmlp_1_rw2_224.sw_in12k_ft_in1k
   * - tu-coatnet_rmlp_1_rw_224.sw_in1k
   * - tu-coatnet_rmlp_2_rw_224.sw_in1k
   * - tu-coatnet_rmlp_2_rw_224.sw_in12k
   * - tu-coatnet_rmlp_2_rw_224.sw_in12k_ft_in1k
   * - tu-coatnet_rmlp_2_rw_384.sw_in12k_ft_in1k
   * - tu-coatnet_rmlp_nano_rw_224.sw_in1k
   * - tu-deit3_base_patch16_224.fb_in1k
   * - tu-deit3_base_patch16_224.fb_in22k_ft_in1k
   * - tu-deit3_base_patch16_384.fb_in1k
   * - tu-deit3_base_patch16_384.fb_in22k_ft_in1k
   * - tu-deit3_large_patch16_224.fb_in1k
   * - tu-deit3_large_patch16_224.fb_in22k_ft_in1k
   * - tu-deit3_large_patch16_384.fb_in1k
   * - tu-deit3_large_patch16_384.fb_in22k_ft_in1k
   * - tu-deit3_medium_patch16_224.fb_in1k
   * - tu-deit3_medium_patch16_224.fb_in22k_ft_in1k
   * - tu-deit3_small_patch16_224.fb_in1k
   * - tu-deit3_small_patch16_224.fb_in22k_ft_in1k
   * - tu-deit3_small_patch16_384.fb_in1k
   * - tu-deit3_small_patch16_384.fb_in22k_ft_in1k
   * - tu-deit_base_distilled_patch16_224.fb_in1k
   * - tu-deit_base_distilled_patch16_384.fb_in1k
   * - tu-deit_base_patch16_224.fb_in1k
   * - tu-deit_base_patch16_384.fb_in1k
   * - tu-deit_small_distilled_patch16_224.fb_in1k
   * - tu-deit_small_patch16_224.fb_in1k
   * - tu-deit_tiny_distilled_patch16_224.fb_in1k
   * - tu-deit_tiny_patch16_224.fb_in1k
   * - tu-regnety_160.deit_in1k
   * - tu-twins_pcpvt_base.in1k
   * - tu-twins_pcpvt_large.in1k
   * - tu-twins_pcpvt_small.in1k
   * - tu-twins_svt_base.in1k
   * - tu-twins_svt_large.in1k
   * - tu-twins_svt_small.in1k
   * - tu-xcit_large_24_p8_224.fb_dist_in1k
   * - tu-xcit_large_24_p8_224.fb_in1k
   * - tu-xcit_large_24_p8_384.fb_dist_in1k
   * - tu-xcit_large_24_p16_224.fb_dist_in1k
   * - tu-xcit_large_24_p16_224.fb_in1k
   * - tu-xcit_large_24_p16_384.fb_dist_in1k
   * - tu-xcit_medium_24_p8_224.fb_dist_in1k
   * - tu-xcit_medium_24_p8_224.fb_in1k
   * - tu-xcit_medium_24_p8_384.fb_dist_in1k
   * - tu-xcit_medium_24_p16_224.fb_dist_in1k
   * - tu-xcit_medium_24_p16_224.fb_in1k
   * - tu-xcit_medium_24_p16_384.fb_dist_in1k
   * - tu-xcit_nano_12_p8_224.fb_dist_in1k
   * - tu-xcit_nano_12_p8_224.fb_in1k
   * - tu-xcit_nano_12_p8_384.fb_dist_in1k
   * - tu-xcit_nano_12_p16_224.fb_dist_in1k
   * - tu-xcit_nano_12_p16_224.fb_in1k
   * - tu-xcit_nano_12_p16_384.fb_dist_in1k
   * - tu-xcit_small_12_p8_224.fb_dist_in1k
   * - tu-xcit_small_12_p8_224.fb_in1k
   * - tu-xcit_small_12_p8_384.fb_dist_in1k
   * - tu-xcit_small_12_p16_224.fb_dist_in1k
   * - tu-xcit_small_12_p16_224.fb_in1k
   * - tu-xcit_small_12_p16_384.fb_dist_in1k
   * - tu-xcit_small_24_p8_224.fb_dist_in1k
   * - tu-xcit_small_24_p8_224.fb_in1k
   * - tu-xcit_small_24_p8_384.fb_dist_in1k
   * - tu-xcit_small_24_p16_224.fb_dist_in1k
   * - tu-xcit_small_24_p16_224.fb_in1k
   * - tu-xcit_small_24_p16_384.fb_dist_in1k
   * - tu-xcit_tiny_12_p8_224.fb_dist_in1k
   * - tu-xcit_tiny_12_p8_224.fb_in1k
   * - tu-xcit_tiny_12_p8_384.fb_dist_in1k
   * - tu-xcit_tiny_12_p16_224.fb_dist_in1k
   * - tu-xcit_tiny_12_p16_224.fb_in1k
   * - tu-xcit_tiny_12_p16_384.fb_dist_in1k
   * - tu-xcit_tiny_24_p8_224.fb_dist_in1k
   * - tu-xcit_tiny_24_p8_224.fb_in1k
   * - tu-xcit_tiny_24_p8_384.fb_dist_in1k
   * - tu-xcit_tiny_24_p16_224.fb_dist_in1k
   * - tu-xcit_tiny_24_p16_224.fb_in1k
   * - tu-xcit_tiny_24_p16_384.fb_dist_in1k


================================================
FILE: docs/encoders_timm.rst
================================================
🎯 Timm Encoders
================

Pytorch Image Models (a.k.a. timm) has a lot of pretrained models and interface which allows using these models as encoders in smp,
however, not all models are supported

 - not all transformer models have ``features_only`` functionality implemented that is required for encoder
 - some models have inappropriate strides

Below is a table of suitable encoders (for DeepLabV3, DeepLabV3+, and PAN dilation support is needed also)

Total number of encoders: 812 (593+219)

.. note::

    To use following encoders you have to add prefix ``tu-``, e.g. ``tu-adv_inception_v3``

Traditional-Style
~~~~~~~~~~~~~~~~~

These models typically produce feature maps at the following downsampling scales relative to the input resolution: 1/2, 1/4, 1/8, 1/16, and 1/32

+----------------------------------+------------------+
| Encoder name                     | Support dilation |
+==================================+==================+
| bat_resnext26ts                  |        ✅        |
+----------------------------------+------------------+
| botnet26t_256                    |        ✅        |
+----------------------------------+------------------+
| botnet50ts_256                   |        ✅        |
+----------------------------------+------------------+
| coatnet_0_224                    |                  |
+----------------------------------+------------------+
| coatnet_0_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_1_224                    |                  |
+----------------------------------+------------------+
| coatnet_1_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_2_224                    |                  |
+----------------------------------+------------------+
| coatnet_2_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_3_224                    |                  |
+----------------------------------+------------------+
| coatnet_3_rw_224                 |                  |
+----------------------------------+------------------+
| coatnet_4_224                    |                  |
+----------------------------------+------------------+
| coatnet_5_224                    |                  |
+----------------------------------+------------------+
| coatnet_bn_0_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_nano_cc_224              |                  |
+----------------------------------+------------------+
| coatnet_nano_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_pico_rw_224              |                  |
+----------------------------------+------------------+
| coatnet_rmlp_0_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_1_rw2_224           |                  |
+----------------------------------+------------------+
| coatnet_rmlp_1_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_2_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_2_rw_384            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_3_rw_224            |                  |
+----------------------------------+------------------+
| coatnet_rmlp_nano_rw_224         |                  |
+----------------------------------+------------------+
| coatnext_nano_rw_224             |                  |
+----------------------------------+------------------+
| cs3darknet_focus_l               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_m               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_s               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_focus_x               |        ✅        |
+----------------------------------+------------------+
| cs3darknet_l                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_m                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_s                     |        ✅        |
+----------------------------------+------------------+
| cs3darknet_x                     |        ✅        |
+----------------------------------+------------------+
| cs3edgenet_x                     |        ✅        |
+----------------------------------+------------------+
| cs3se_edgenet_x                  |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_l                   |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_x                   |        ✅        |
+----------------------------------+------------------+
| cs3sedarknet_xdw                 |        ✅        |
+----------------------------------+------------------+
| cspdarknet53                     |        ✅        |
+----------------------------------+------------------+
| cspresnet50                      |        ✅        |
+----------------------------------+------------------+
| cspresnet50d                     |        ✅        |
+----------------------------------+------------------+
| cspresnet50w                     |        ✅        |
+----------------------------------+------------------+
| cspresnext50                     |        ✅        |
+----------------------------------+------------------+
| darknet17                        |        ✅        |
+----------------------------------+------------------+
| darknet21                        |        ✅        |
+----------------------------------+------------------+
| darknet53                        |        ✅        |
+----------------------------------+------------------+
| darknetaa53                      |        ✅        |
+----------------------------------+------------------+
| densenet121                      |                  |
+----------------------------------+------------------+
| densenet161                      |                  |
+----------------------------------+------------------+
| densenet169                      |                  |
+----------------------------------+------------------+
| densenet201                      |                  |
+----------------------------------+------------------+
| densenet264d                     |                  |
+----------------------------------+------------------+
| densenetblur121d                 |                  |
+----------------------------------+------------------+
| dla34                            |                  |
+----------------------------------+------------------+
| dla46_c                          |                  |
+----------------------------------+------------------+
| dla46x_c                         |                  |
+----------------------------------+------------------+
| dla60                            |                  |
+----------------------------------+------------------+
| dla60_res2net                    |                  |
+----------------------------------+------------------+
| dla60_res2next                   |                  |
+----------------------------------+------------------+
| dla60x                           |                  |
+----------------------------------+------------------+
| dla60x_c                         |                  |
+----------------------------------+------------------+
| dla102                           |                  |
+----------------------------------+------------------+
| dla102x                          |                  |
+----------------------------------+------------------+
| dla102x2                         |                  |
+----------------------------------+------------------+
| dla169                           |                  |
+----------------------------------+------------------+
| dm_nfnet_f0                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f1                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f2                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f3                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f4                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f5                      |        ✅        |
+----------------------------------+------------------+
| dm_nfnet_f6                      |        ✅        |
+----------------------------------+------------------+
| dpn48b                           |                  |
+----------------------------------+------------------+
| dpn68                            |                  |
+----------------------------------+------------------+
| dpn68b                           |                  |
+----------------------------------+------------------+
| dpn92                            |                  |
+----------------------------------+------------------+
| dpn98                            |                  |
+----------------------------------+------------------+
| dpn107                           |                  |
+----------------------------------+------------------+
| dpn131                           |                  |
+----------------------------------+------------------+
| eca_botnext26ts_256              |        ✅        |
+----------------------------------+------------------+
| eca_halonext26ts                 |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l0                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l1                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l2                     |        ✅        |
+----------------------------------+------------------+
| eca_nfnet_l3                     |        ✅        |
+----------------------------------+------------------+
| eca_resnet33ts                   |        ✅        |
+----------------------------------+------------------+
| eca_resnext26ts                  |        ✅        |
+----------------------------------+------------------+
| eca_vovnet39b                    |                  |
+----------------------------------+------------------+
| ecaresnet26t                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet50d                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet50d_pruned              |        ✅        |
+----------------------------------+------------------+
| ecaresnet50t                     |        ✅        |
+----------------------------------+------------------+
| ecaresnet101d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnet101d_pruned             |        ✅        |
+----------------------------------+------------------+
| ecaresnet200d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnet269d                    |        ✅        |
+----------------------------------+------------------+
| ecaresnetlight                   |        ✅        |
+----------------------------------+------------------+
| ecaresnext26t_32x4d              |        ✅        |
+----------------------------------+------------------+
| ecaresnext50t_32x4d              |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_g8_gn            |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_g16_evos         |        ✅        |
+----------------------------------+------------------+
| efficientnet_b0_gn               |        ✅        |
+----------------------------------+------------------+
| efficientnet_b1                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b1_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b2                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b2_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_g8_gn            |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_gn               |        ✅        |
+----------------------------------+------------------+
| efficientnet_b3_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_b4                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b5                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b6                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b7                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_b8                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_blur_b0             |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b0_4e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b0_8e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_cc_b1_8e            |        ✅        |
+----------------------------------+------------------+
| efficientnet_el                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_el_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_em                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_es                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_es_pruned           |        ✅        |
+----------------------------------+------------------+
| efficientnet_l2                  |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite0               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite1               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite2               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite3               |        ✅        |
+----------------------------------+------------------+
| efficientnet_lite4               |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_l                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_m                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_m              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_s              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_rw_t              |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_s                 |        ✅        |
+----------------------------------+------------------+
| efficientnetv2_xl                |        ✅        |
+----------------------------------+------------------+
| ese_vovnet19b_dw                 |                  |
+----------------------------------+------------------+
| ese_vovnet19b_slim               |                  |
+----------------------------------+------------------+
| ese_vovnet19b_slim_dw            |                  |
+----------------------------------+------------------+
| ese_vovnet39b                    |                  |
+----------------------------------+------------------+
| ese_vovnet39b_evos               |                  |
+----------------------------------+------------------+
| ese_vovnet57b                    |                  |
+----------------------------------+------------------+
| ese_vovnet99b                    |                  |
+----------------------------------+------------------+
| fbnetc_100                       |        ✅        |
+----------------------------------+------------------+
| fbnetv3_b                        |        ✅        |
+----------------------------------+------------------+
| fbnetv3_d                        |        ✅        |
+----------------------------------+------------------+
| fbnetv3_g                        |        ✅        |
+----------------------------------+------------------+
| gc_efficientnetv2_rw_t           |        ✅        |
+----------------------------------+------------------+
| gcresnet33ts                     |        ✅        |
+----------------------------------+------------------+
| gcresnet50t                      |        ✅        |
+----------------------------------+------------------+
| gcresnext26ts                    |        ✅        |
+----------------------------------+------------------+
| gcresnext50ts                    |        ✅        |
+----------------------------------+------------------+
| gernet_l                         |        ✅        |
+----------------------------------+------------------+
| gernet_m                         |        ✅        |
+----------------------------------+------------------+
| gernet_s                         |        ✅        |
+----------------------------------+------------------+
| ghostnet_050                     |                  |
+----------------------------------+------------------+
| ghostnet_100                     |                  |
+----------------------------------+------------------+
| ghostnet_130                     |                  |
+----------------------------------+------------------+
| ghostnetv2_100                   |                  |
+----------------------------------+------------------+
| ghostnetv2_130                   |                  |
+----------------------------------+------------------+
| ghostnetv2_160                   |                  |
+----------------------------------+------------------+
| halo2botnet50ts_256              |        ✅        |
+----------------------------------+------------------+
| halonet26t                       |        ✅        |
+----------------------------------+------------------+
| halonet50ts                      |        ✅        |
+----------------------------------+------------------+
| halonet_h1                       |        ✅        |
+----------------------------------+------------------+
| haloregnetz_b                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_a                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_b                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_c                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_d                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_e                    |        ✅        |
+----------------------------------+------------------+
| hardcorenas_f                    |        ✅        |
+----------------------------------+------------------+
| hrnet_w18                        |                  |
+----------------------------------+------------------+
| hrnet_w18_small                  |                  |
+----------------------------------+------------------+
| hrnet_w18_small_v2               |                  |
+----------------------------------+------------------+
| hrnet_w18_ssld                   |                  |
+----------------------------------+------------------+
| hrnet_w30                        |                  |
+----------------------------------+------------------+
| hrnet_w32                        |                  |
+----------------------------------+------------------+
| hrnet_w40                        |                  |
+----------------------------------+------------------+
| hrnet_w44                        |                  |
+----------------------------------+------------------+
| hrnet_w48                        |                  |
+----------------------------------+------------------+
| hrnet_w48_ssld                   |                  |
+----------------------------------+------------------+
| hrnet_w64                        |                  |
+----------------------------------+------------------+
| lambda_resnet26rpt_256           |        ✅        |
+----------------------------------+------------------+
| lambda_resnet26t                 |        ✅        |
+----------------------------------+------------------+
| lambda_resnet50ts                |        ✅        |
+----------------------------------+------------------+
| lamhalobotnet50ts_256            |        ✅        |
+----------------------------------+------------------+
| lcnet_035                        |        ✅        |
+----------------------------------+------------------+
| lcnet_050                        |        ✅        |
+----------------------------------+------------------+
| lcnet_075                        |        ✅        |
+----------------------------------+------------------+
| lcnet_100                        |        ✅        |
+----------------------------------+------------------+
| lcnet_150                        |        ✅        |
+----------------------------------+------------------+
| legacy_senet154                  |                  |
+----------------------------------+------------------+
| legacy_seresnet18                |                  |
+----------------------------------+------------------+
| legacy_seresnet34                |                  |
+----------------------------------+------------------+
| legacy_seresnet50                |                  |
+----------------------------------+------------------+
| legacy_seresnet101               |                  |
+----------------------------------+------------------+
| legacy_seresnet152               |                  |
+----------------------------------+------------------+
| legacy_seresnext26_32x4d         |                  |
+----------------------------------+------------------+
| legacy_seresnext50_32x4d         |                  |
+----------------------------------+------------------+
| legacy_seresnext101_32x4d        |                  |
+----------------------------------+------------------+
| maxvit_base_tf_224               |                  |
+----------------------------------+------------------+
| maxvit_base_tf_384               |                  |
+----------------------------------+------------------+
| maxvit_base_tf_512               |                  |
+----------------------------------+------------------+
| maxvit_large_tf_224              |                  |
+----------------------------------+------------------+
| maxvit_large_tf_384              |                  |
+----------------------------------+------------------+
| maxvit_large_tf_512              |                  |
+----------------------------------+------------------+
| maxvit_nano_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_pico_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_rmlp_base_rw_224          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_base_rw_384          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_nano_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_pico_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_rmlp_small_rw_224         |                  |
+----------------------------------+------------------+
| maxvit_rmlp_small_rw_256         |                  |
+----------------------------------+------------------+
| maxvit_rmlp_tiny_rw_256          |                  |
+----------------------------------+------------------+
| maxvit_small_tf_224              |                  |
+----------------------------------+------------------+
| maxvit_small_tf_384              |                  |
+----------------------------------+------------------+
| maxvit_small_tf_512              |                  |
+----------------------------------+------------------+
| maxvit_tiny_pm_256               |                  |
+----------------------------------+------------------+
| maxvit_tiny_rw_224               |                  |
+----------------------------------+------------------+
| maxvit_tiny_rw_256               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_224               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_384               |                  |
+----------------------------------+------------------+
| maxvit_tiny_tf_512               |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_224             |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_384             |                  |
+----------------------------------+------------------+
| maxvit_xlarge_tf_512             |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_nano_rw_256         |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_small_rw_256        |                  |
+----------------------------------+------------------+
| maxxvit_rmlp_tiny_rw_256         |                  |
+----------------------------------+------------------+
| maxxvitv2_nano_rw_256            |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_base_rw_224       |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_base_rw_384       |                  |
+----------------------------------+------------------+
| maxxvitv2_rmlp_large_rw_224      |                  |
+----------------------------------+------------------+
| mixnet_l                         |        ✅        |
+----------------------------------+------------------+
| mixnet_m                         |        ✅        |
+----------------------------------+------------------+
| mixnet_s                         |        ✅        |
+----------------------------------+------------------+
| mixnet_xl                        |        ✅        |
+----------------------------------+------------------+
| mixnet_xxl                       |        ✅        |
+----------------------------------+------------------+
| mnasnet_050                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_075                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_100                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_140                      |        ✅        |
+----------------------------------+------------------+
| mnasnet_small                    |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_100            |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_l           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_m           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_s           |        ✅        |
+----------------------------------+------------------+
| mobilenet_edgetpu_v2_xs          |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_100                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_100h                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv1_125                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_035                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_050                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_075                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_100                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_110d                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_120d                 |        ✅        |
+----------------------------------+------------------+
| mobilenetv2_140                  |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_075            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_100            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_large_150d           |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_rw                   |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_050            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_075            |        ✅        |
+----------------------------------+------------------+
| mobilenetv3_small_100            |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_aa_large        |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_aa_medium       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_blur_medium     |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_large           |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_medium          |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small           |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small_035       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_conv_small_050       |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_large         |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_large_075     |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_medium        |        ✅        |
+----------------------------------+------------------+
| mobilenetv4_hybrid_medium_075    |        ✅        |
+----------------------------------+------------------+
| mobileone_s0                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s1                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s2                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s3                     |        ✅        |
+----------------------------------+------------------+
| mobileone_s4                     |        ✅        |
+----------------------------------+------------------+
| mobilevit_s                      |        ✅        |
+----------------------------------+------------------+
| mobilevit_xs                     |        ✅        |
+----------------------------------+------------------+
| mobilevit_xxs                    |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_050                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_075                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_100                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_125                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_150                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_175                  |        ✅        |
+----------------------------------+------------------+
| mobilevitv2_200                  |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet26                   |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet50                   |        ✅        |
+----------------------------------+------------------+
| nf_ecaresnet101                  |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b0                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b1                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b2                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b3                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b4                     |        ✅        |
+----------------------------------+------------------+
| nf_regnet_b5                     |        ✅        |
+----------------------------------+------------------+
| nf_resnet26                      |        ✅        |
+----------------------------------+------------------+
| nf_resnet50                      |        ✅        |
+----------------------------------+------------------+
| nf_resnet101                     |        ✅        |
+----------------------------------+------------------+
| nf_seresnet26                    |        ✅        |
+----------------------------------+------------------+
| nf_seresnet50                    |        ✅        |
+----------------------------------+------------------+
| nf_seresnet101                   |        ✅        |
+----------------------------------+------------------+
| nfnet_f0                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f1                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f2                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f3                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f4                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f5                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f6                         |        ✅        |
+----------------------------------+------------------+
| nfnet_f7                         |        ✅        |
+----------------------------------+------------------+
| nfnet_l0                         |        ✅        |
+----------------------------------+------------------+
| regnetv_040                      |        ✅        |
+----------------------------------+------------------+
| regnetv_064                      |        ✅        |
+----------------------------------+------------------+
| regnetx_002                      |        ✅        |
+----------------------------------+------------------+
| regnetx_004                      |        ✅        |
+----------------------------------+------------------+
| regnetx_004_tv                   |        ✅        |
+----------------------------------+------------------+
| regnetx_006                      |        ✅        |
+----------------------------------+------------------+
| regnetx_008                      |        ✅        |
+----------------------------------+------------------+
| regnetx_016                      |        ✅        |
+----------------------------------+------------------+
| regnetx_032                      |        ✅        |
+----------------------------------+------------------+
| regnetx_040                      |        ✅        |
+----------------------------------+------------------+
| regnetx_064                      |        ✅        |
+----------------------------------+------------------+
| regnetx_080                      |        ✅        |
+----------------------------------+------------------+
| regnetx_120                      |        ✅        |
+----------------------------------+------------------+
| regnetx_160                      |        ✅        |
+----------------------------------+------------------+
| regnetx_320                      |        ✅        |
+----------------------------------+------------------+
| regnety_002                      |        ✅        |
+----------------------------------+------------------+
| regnety_004                      |        ✅        |
+----------------------------------+------------------+
| regnety_006                      |        ✅        |
+----------------------------------+------------------+
| regnety_008                      |        ✅        |
+----------------------------------+------------------+
| regnety_008_tv                   |        ✅        |
+----------------------------------+------------------+
| regnety_016                      |        ✅        |
+----------------------------------+------------------+
| regnety_032                      |        ✅        |
+----------------------------------+------------------+
| regnety_040                      |        ✅        |
+----------------------------------+------------------+
| regnety_040_sgn                  |        ✅        |
+----------------------------------+------------------+
| regnety_064                      |        ✅        |
+----------------------------------+------------------+
| regnety_080                      |        ✅        |
+----------------------------------+------------------+
| regnety_080_tv                   |        ✅        |
+----------------------------------+------------------+
| regnety_120                      |        ✅        |
+----------------------------------+------------------+
| regnety_160                      |        ✅        |
+----------------------------------+------------------+
| regnety_320                      |        ✅        |
+----------------------------------+------------------+
| regnety_640                      |        ✅        |
+----------------------------------+------------------+
| regnety_1280                     |        ✅        |
+----------------------------------+------------------+
| regnety_2560                     |        ✅        |
+----------------------------------+------------------+
| regnetz_005                      |        ✅        |
+----------------------------------+------------------+
| regnetz_040                      |        ✅        |
+----------------------------------+------------------+
| regnetz_040_h                    |        ✅        |
+----------------------------------+------------------+
| regnetz_b16                      |        ✅        |
+----------------------------------+------------------+
| regnetz_b16_evos                 |        ✅        |
+----------------------------------+------------------+
| regnetz_c16                      |        ✅        |
+----------------------------------+------------------+
| regnetz_c16_evos                 |        ✅        |
+----------------------------------+------------------+
| regnetz_d8                       |        ✅        |
+----------------------------------+------------------+
| regnetz_d8_evos                  |        ✅        |
+----------------------------------+------------------+
| regnetz_d32                      |        ✅        |
+----------------------------------+------------------+
| regnetz_e8                       |        ✅        |
+----------------------------------+------------------+
| repghostnet_050                  |                  |
+----------------------------------+------------------+
| repghostnet_058                  |                  |
+----------------------------------+------------------+
| repghostnet_080                  |                  |
+----------------------------------+------------------+
| repghostnet_100                  |                  |
+----------------------------------+------------------+
| repghostnet_111                  |                  |
+----------------------------------+------------------+
| repghostnet_130                  |                  |
+----------------------------------+------------------+
| repghostnet_150                  |                  |
+----------------------------------+------------------+
| repghostnet_200                  |                  |
+----------------------------------+------------------+
| repvgg_a0                        |        ✅        |
+----------------------------------+------------------+
| repvgg_a1                        |        ✅        |
+----------------------------------+------------------+
| repvgg_a2                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b0                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b1                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b1g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_b2                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b2g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_b3                        |        ✅        |
+----------------------------------+------------------+
| repvgg_b3g4                      |        ✅        |
+----------------------------------+------------------+
| repvgg_d2se                      |        ✅        |
+----------------------------------+------------------+
| res2net50_14w_8s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_4s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_6s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_26w_8s                 |        ✅        |
+----------------------------------+------------------+
| res2net50_48w_2s                 |        ✅        |
+----------------------------------+------------------+
| res2net50d                       |        ✅        |
+----------------------------------+------------------+
| res2net101_26w_4s                |        ✅        |
+----------------------------------+------------------+
| res2net101d                      |        ✅        |
+----------------------------------+------------------+
| res2next50                       |        ✅        |
+----------------------------------+------------------+
| resnest14d                       |        ✅        |
+----------------------------------+------------------+
| resnest26d                       |        ✅        |
+----------------------------------+------------------+
| resnest50d                       |        ✅        |
+----------------------------------+------------------+
| resnest50d_1s4x24d               |        ✅        |
+----------------------------------+------------------+
| resnest50d_4s2x40d               |        ✅        |
+----------------------------------+------------------+
| resnest101e                      |        ✅        |
+----------------------------------+------------------+
| resnest200e                      |        ✅        |
+----------------------------------+------------------+
| resnest269e                      |        ✅        |
+----------------------------------+------------------+
| resnet10t                        |        ✅        |
+----------------------------------+------------------+
| resnet14t                        |        ✅        |
+----------------------------------+------------------+
| resnet18                         |        ✅        |
+----------------------------------+------------------+
| resnet18d                        |        ✅        |
+----------------------------------+------------------+
| resnet26                         |        ✅        |
+----------------------------------+------------------+
| resnet26d                        |        ✅        |
+----------------------------------+------------------+
| resnet26t                        |        ✅        |
+----------------------------------+------------------+
| resnet32ts                       |        ✅        |
+----------------------------------+------------------+
| resnet33ts                       |        ✅        |
+----------------------------------+------------------+
| resnet34                         |        ✅        |
+----------------------------------+------------------+
| resnet34d                        |        ✅        |
+----------------------------------+------------------+
| resnet50                         |        ✅        |
+----------------------------------+------------------+
| resnet50_clip                    |        ✅        |
+----------------------------------+------------------+
| resnet50_clip_gap                |        ✅        |
+----------------------------------+------------------+
| resnet50_gn                      |        ✅        |
+----------------------------------+------------------+
| resnet50_mlp                     |        ✅        |
+----------------------------------+------------------+
| resnet50c                        |        ✅        |
+----------------------------------+------------------+
| resnet50d                        |        ✅        |
+----------------------------------+------------------+
| resnet50s                        |        ✅        |
+----------------------------------+------------------+
| resnet50t                        |        ✅        |
+----------------------------------+------------------+
| resnet50x4_clip                  |        ✅        |
+----------------------------------+------------------+
| resnet50x4_clip_gap              |        ✅        |
+----------------------------------+------------------+
| resnet50x16_clip                 |        ✅        |
+----------------------------------+------------------+
| resnet50x16_clip_gap             |        ✅        |
+----------------------------------+------------------+
| resnet50x64_clip                 |        ✅        |
+----------------------------------+------------------+
| resnet50x64_clip_gap             |        ✅        |
+----------------------------------+------------------+
| resnet51q                        |        ✅        |
+----------------------------------+------------------+
| resnet61q                        |        ✅        |
+----------------------------------+------------------+
| resnet101                        |        ✅        |
+----------------------------------+------------------+
| resnet101_clip                   |        ✅        |
+----------------------------------+------------------+
| resnet101_clip_gap               |        ✅        |
+----------------------------------+------------------+
| resnet101c                       |        ✅        |
+----------------------------------+------------------+
| resnet101d                       |        ✅        |
+----------------------------------+------------------+
| resnet101s                       |        ✅        |
+----------------------------------+------------------+
| resnet152                        |        ✅        |
+----------------------------------+------------------+
| resnet152c                       |        ✅        |
+----------------------------------+------------------+
| resnet152d                       |        ✅        |
+----------------------------------+------------------+
| resnet152s                       |        ✅        |
+----------------------------------+------------------+
| resnet200                        |        ✅        |
+----------------------------------+------------------+
| resnet200d                       |        ✅        |
+----------------------------------+------------------+
| resnetaa34d                      |        ✅        |
+----------------------------------+------------------+
| resnetaa50                       |        ✅        |
+----------------------------------+------------------+
| resnetaa50d                      |        ✅        |
+----------------------------------+------------------+
| resnetaa101d                     |        ✅        |
+----------------------------------+------------------+
| resnetblur18                     |        ✅        |
+----------------------------------+------------------+
| resnetblur50                     |        ✅        |
+----------------------------------+------------------+
| resnetblur50d                    |        ✅        |
+----------------------------------+------------------+
| resnetblur101d                   |        ✅        |
+----------------------------------+------------------+
| resnetrs50                       |        ✅        |
+----------------------------------+------------------+
| resnetrs101                      |        ✅        |
+----------------------------------+------------------+
| resnetrs152                      |        ✅        |
+----------------------------------+------------------+
| resnetrs200                      |        ✅        |
+----------------------------------+------------------+
| resnetrs270                      |        ✅        |
+----------------------------------+------------------+
| resnetrs350                      |        ✅        |
+----------------------------------+------------------+
| resnetrs420                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_18                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_18d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_34                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_34d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50                      |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_evos                |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_frn                 |        ✅        |
+----------------------------------+------------------+
| resnetv2_50d_gn                  |        ✅        |
+----------------------------------+------------------+
| resnetv2_50t                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_50x1_bit                |        ✅        |
+----------------------------------+------------------+
| resnetv2_50x3_bit                |        ✅        |
+----------------------------------+------------------+
| resnetv2_101                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_101d                    |        ✅        |
+----------------------------------+------------------+
| resnetv2_101x1_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_101x3_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_152                     |        ✅        |
+----------------------------------+------------------+
| resnetv2_152d                    |        ✅        |
+----------------------------------+------------------+
| resnetv2_152x2_bit               |        ✅        |
+----------------------------------+------------------+
| resnetv2_152x4_bit               |        ✅        |
+----------------------------------+------------------+
| resnext26ts                      |        ✅        |
+----------------------------------+------------------+
| resnext50_32x4d                  |        ✅        |
+----------------------------------+------------------+
| resnext50d_32x4d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x4d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x8d                 |        ✅        |
+----------------------------------+------------------+
| resnext101_32x16d                |        ✅        |
+----------------------------------+------------------+
| resnext101_32x32d                |        ✅        |
+----------------------------------+------------------+
| resnext101_64x4d                 |        ✅        |
+----------------------------------+------------------+
| rexnet_100                       |        ✅        |
+----------------------------------+------------------+
| rexnet_130                       |        ✅        |
+----------------------------------+------------------+
| rexnet_150                       |        ✅        |
+----------------------------------+------------------+
| rexnet_200                       |        ✅        |
+----------------------------------+------------------+
| rexnet_300                       |        ✅        |
+----------------------------------+------------------+
| rexnetr_100                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_130                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_150                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_200                      |        ✅        |
+----------------------------------+------------------+
| rexnetr_300                      |        ✅        |
+----------------------------------+------------------+
| sebotnet33ts_256                 |        ✅        |
+----------------------------------+------------------+
| sedarknet21                      |        ✅        |
+----------------------------------+------------------+
| sehalonet33ts                    |        ✅        |
+----------------------------------+------------------+
| selecsls42                       |                  |
+----------------------------------+------------------+
| selecsls42b                      |                  |
+----------------------------------+------------------+
| selecsls60                       |                  |
+----------------------------------+------------------+
| selecsls60b                      |                  |
+----------------------------------+------------------+
| selecsls84                       |                  |
+----------------------------------+------------------+
| semnasnet_050                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_075                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_100                    |        ✅        |
+----------------------------------+------------------+
| semnasnet_140                    |        ✅        |
+----------------------------------+------------------+
| senet154                         |        ✅        |
+----------------------------------+------------------+
| seresnet18                       |        ✅        |
+----------------------------------+------------------+
| seresnet33ts                     |        ✅        |
+----------------------------------+------------------+
| seresnet34                       |        ✅        |
+----------------------------------+------------------+
| seresnet50                       |        ✅        |
+----------------------------------+------------------+
| seresnet50t                      |        ✅        |
+----------------------------------+------------------+
| seresnet101                      |        ✅        |
+----------------------------------+------------------+
| seresnet152                      |        ✅        |
+----------------------------------+------------------+
| seresnet152d                     |        ✅        |
+----------------------------------+------------------+
| seresnet200d                     |        ✅        |
+----------------------------------+------------------+
| seresnet269d                     |        ✅        |
+----------------------------------+------------------+
| seresnetaa50d                    |        ✅        |
+----------------------------------+------------------+
| seresnext26d_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext26t_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext26ts                    |        ✅        |
+----------------------------------+------------------+
| seresnext50_32x4d                |        ✅        |
+----------------------------------+------------------+
| seresnext101_32x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext101_32x8d               |        ✅        |
+----------------------------------+------------------+
| seresnext101_64x4d               |        ✅        |
+----------------------------------+------------------+
| seresnext101d_32x8d              |        ✅        |
+----------------------------------+------------------+
| seresnextaa101d_32x8d            |        ✅        |
+----------------------------------+------------------+
| seresnextaa201d_32x8d            |        ✅        |
+----------------------------------+------------------+
| skresnet18                       |        ✅        |
+----------------------------------+------------------+
| skresnet34                       |        ✅        |
+----------------------------------+------------------+
| skresnet50                       |        ✅        |
+----------------------------------+------------------+
| skresnet50d                      |        ✅        |
+----------------------------------+------------------+
| skresnext50_32x4d                |        ✅        |
+----------------------------------+------------------+
| spnasnet_100                     |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b0               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b1               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b2               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b3               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b4               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b5               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b6               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b7               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_b8               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b0_4e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b0_8e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_cc_b1_8e         |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_el               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_em               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_es               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_l2               |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite0            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite1            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite2            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite3            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnet_lite4            |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b0             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b1             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b2             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_b3             |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_l              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_m              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_s              |        ✅        |
+----------------------------------+------------------+
| tf_efficientnetv2_xl             |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_l                      |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_m                      |        ✅        |
+----------------------------------+------------------+
| tf_mixnet_s                      |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_075         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_100         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_large_minimal_100 |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_075         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_100         |        ✅        |
+----------------------------------+------------------+
| tf_mobilenetv3_small_minimal_100 |        ✅        |
+----------------------------------+------------------+
| tinynet_a                        |        ✅        |
+----------------------------------+------------------+
| tinynet_b                        |        ✅        |
+----------------------------------+------------------+
| tinynet_c                        |        ✅        |
+----------------------------------+------------------+
| tinynet_d                        |        ✅        |
+----------------------------------+------------------+
| tinynet_e                        |        ✅        |
+----------------------------------+------------------+
| vgg11                            |                  |
+----------------------------------+------------------+
| vgg11_bn                         |                  |
+----------------------------------+------------------+
| vgg13                            |                  |
+----------------------------------+------------------+
| vgg13_bn                         |                  |
+----------------------------------+------------------+
| vgg16                            |                  |
+----------------------------------+------------------+
| vgg16_bn                         |                  |
+----------------------------------+------------------+
| vgg19                            |                  |
+----------------------------------+------------------+
| vgg19_bn                         |                  |
+----------------------------------+------------------+
| vovnet39a                        |                  |
+----------------------------------+------------------+
| vovnet57a                        |                  |
+----------------------------------+------------------+
| wide_resnet50_2                  |        ✅        |
+----------------------------------+------------------+
| wide_resnet101_2                 |        ✅        |
+----------------------------------+------------------+
| xception41                       |        ✅        |
+----------------------------------+------------------+
| xception41p                      |        ✅        |
+----------------------------------+------------------+
| xception65                       |        ✅        |
+----------------------------------+------------------+
| xception65p                      |        ✅        |
+----------------------------------+------------------+
| xception71                       |        ✅        |
+----------------------------------+------------------+

Transformer-Style
~~~~~~~~~~~~~~~~~

Transformer-style models (e.g., Swin Transformer, ConvNeXt) typically produce feature maps starting at a 1/4 scale, followed by 1/8, 1/16, and 1/32 scales

+------------------------------------+------------------+
| Encoder name                       | Support dilation |
+====================================+==================+
| caformer_b36                       |                  |
+------------------------------------+------------------+
| caformer_m36                       |                  |
+------------------------------------+------------------+
| caformer_s18                       |                  |
+------------------------------------+------------------+
| caformer_s36                       |                  |
+------------------------------------+------------------+
| convformer_b36                     |                  |
+------------------------------------+------------------+
| convformer_m36                     |                  |
+------------------------------------+------------------+
| convformer_s18                     |                  |
+------------------------------------+------------------+
| convformer_s36                     |                  |
+------------------------------------+------------------+
| convnext_atto                      |        ✅        |
+------------------------------------+------------------+
| convnext_atto_ols                  |        ✅        |
+------------------------------------+------------------+
| convnext_atto_rms                  |        ✅        |
+------------------------------------+------------------+
| convnext_base                      |        ✅        |
+------------------------------------+------------------+
| convnext_femto                     |        ✅        |
+------------------------------------+------------------+
| convnext_femto_ols                 |        ✅        |
+------------------------------------+------------------+
| convnext_large                     |        ✅        |
+------------------------------------+------------------+
| convnext_large_mlp                 |        ✅        |
+------------------------------------+------------------+
| convnext_nano                      |        ✅        |
+------------------------------------+------------------+
| convnext_nano_ols                  |        ✅        |
+------------------------------------+------------------+
| convnext_pico                      |        ✅        |
+------------------------------------+------------------+
| convnext_pico_ols                  |        ✅        |
+------------------------------------+------------------+
| convnext_small                     |        ✅        |
+------------------------------------+------------------+
| convnext_tiny                      |        ✅        |
+------------------------------------+------------------+
| convnext_tiny_hnf                  |        ✅        |
+------------------------------------+------------------+
| convnext_xlarge                    |        ✅        |
+------------------------------------+------------------+
| convnext_xxlarge                   |        ✅        |
+------------------------------------+------------------+
| convnext_zepto_rms                 |        ✅        |
+------------------------------------+------------------+
| convnext_zepto_rms_ols             |        ✅        |
+------------------------------------+------------------+
| convnextv2_atto                    |        ✅        |
+------------------------------------+------------------+
| convnextv2_base                    |        ✅        |
+------------------------------------+------------------+
| convnextv2_femto                   |        ✅        |
+------------------------------------+------------------+
| convnextv2_huge                    |        ✅        |
+------------------------------------+------------------+
| convnextv2_large                   |        ✅        |
+------------------------------------+------------------+
| convnextv2_nano                    |        ✅        |
+------------------------------------+------------------+
| convnextv2_pico                    |        ✅        |
+------------------------------------+------------------+
| convnextv2_small                   |        ✅        |
+------------------------------------+------------------+
| convnextv2_tiny                    |        ✅        |
+------------------------------------+------------------+
| davit_base                         |                  |
+------------------------------------+------------------+
| davit_base_fl                      |                  |
+------------------------------------+------------------+
| davit_giant                        |                  |
+------------------------------------+------------------+
| davit_huge                         |                  |
+------------------------------------+------------------+
| davit_huge_fl                      |                  |
+------------------------------------+------------------+
| davit_large                        |                  |
+------------------------------------+------------------+
| davit_small                        |                  |
+------------------------------------+------------------+
| davit_tiny                         |                  |
+------------------------------------+------------------+
| edgenext_base                      |                  |
+------------------------------------+------------------+
| edgenext_small                     |                  |
+------------------------------------+------------------+
| edgenext_small_rw                  |                  |
+------------------------------------+------------------+
| edgenext_x_small                   |                  |
+------------------------------------+------------------+
| edgenext_xx_small                  |                  |
+------------------------------------+------------------+
| efficientformer_l1                 |                  |
+------------------------------------+------------------+
| efficientformer_l3                 |                  |
+------------------------------------+------------------+
| efficientformer_l7                 |                  |
+------------------------------------+------------------+
| efficientformerv2_l                |                  |
+------------------------------------+------------------+
| efficientformerv2_s0               |                  |
+------------------------------------+------------------+
| efficientformerv2_s1               |                  |
+------------------------------------+------------------+
| efficientformerv2_s2               |                  |
+------------------------------------+------------------+
| efficientvit_b0                    |                  |
+------------------------------------+------------------+
| efficientvit_b1                    |                  |
+------------------------------------+------------------+
| efficientvit_b2                    |                  |
+------------------------------------+------------------+
| efficientvit_b3                    |                  |
+------------------------------------+------------------+
| efficientvit_l1                    |                  |
+------------------------------------+------------------+
| efficientvit_l2                    |                  |
+------------------------------------+------------------+
| efficientvit_l3                    |                  |
+------------------------------------+------------------+
| fastvit_ma36                       |                  |
+------------------------------------+------------------+
| fastvit_mci0                       |                  |
+------------------------------------+------------------+
| fastvit_mci1                       |                  |
+------------------------------------+------------------+
| fastvit_mci2                       |                  |
+------------------------------------+------------------+
| fastvit_s12                        |                  |
+------------------------------------+------------------+
| fastvit_sa12                       |                  |
+------------------------------------+------------------+
| fastvit_sa24                       |                  |
+------------------------------------+------------------+
| fastvit_sa36                       |                  |
+------------------------------------+------------------+
| fastvit_t8                         |                  |
+------------------------------------+------------------+
| fastvit_t12                        |                  |
+------------------------------------+------------------+
| focalnet_base_lrf                  |                  |
+------------------------------------+------------------+
| focalnet_base_srf                  |                  |
+------------------------------------+------------------+
| focalnet_huge_fl3                  |                  |
+------------------------------------+------------------+
| focalnet_huge_fl4                  |                  |
+------------------------------------+------------------+
| focalnet_large_fl3                 |                  |
+------------------------------------+------------------+
| focalnet_large_fl4                 |                  |
+------------------------------------+------------------+
| focalnet_small_lrf                 |                  |
+------------------------------------+------------------+
| focalnet_small_srf                 |                  |
+------------------------------------+------------------+
| focalnet_tiny_lrf                  |                  |
+------------------------------------+------------------+
| focalnet_tiny_srf                  |                  |
+------------------------------------+------------------+
| focalnet_xlarge_fl3                |                  |
+------------------------------------+------------------+
| focalnet_xlarge_fl4                |                  |
+------------------------------------+------------------+
| hgnet_base                         |                  |
+------------------------------------+------------------+
| hgnet_small                        |                  |
+------------------------------------+------------------+
| hgnet_tiny                         |                  |
+------------------------------------+------------------+
| hgnetv2_b0                         |                  |
+------------------------------------+------------------+
| hgnetv2_b1                         |                  |
+------------------------------------+------------------+
| hgnetv2_b2                         |                  |
+------------------------------------+------------------+
| hgnetv2_b3                         |                  |
+------------------------------------+------------------+
| hgnetv2_b4                         |                  |
+------------------------------------+------------------+
| hgnetv2_b5                         |                  |
+------------------------------------+------------------+
| hgnetv2_b6                         |                  |
+------------------------------------+------------------+
| hiera_base_224                     |                  |
+------------------------------------+------------------+
| hiera_base_abswin_256              |                  |
+------------------------------------+------------------+
| hiera_base_plus_224                |                  |
+------------------------------------+------------------+
| hiera_huge_224                     |                  |
+------------------------------------+------------------+
| hiera_large_224                    |                  |
+------------------------------------+------------------+
| hiera_small_224                    |                  |
+------------------------------------+------------------+
| hiera_small_abswin_256             |                  |
+------------------------------------+------------------+
| hiera_tiny_224                     |                  |
+------------------------------------+------------------+
| hieradet_small                     |                  |
+------------------------------------+------------------+
| inception_next_base                |                  |
+------------------------------------+------------------+
| inception_next_small               |                  |
+------------------------------------+------------------+
| inception_next_tiny                |                  |
+------------------------------------+------------------+
| mambaout_base                      |                  |
+------------------------------------+------------------+
| mambaout_base_plus_rw              |                  |
+------------------------------------+------------------+
| mambaout_base_short_rw             |                  |
+------------------------------------+------------------+
| mambaout_base_tall_rw              |                  |
+------------------------------------+------------------+
| mambaout_base_wide_rw              |                  |
+------------------------------------+------------------+
| mambaout_femto                     |                  |
+------------------------------------+------------------+
| mambaout_kobe                      |                  |
+------------------------------------+------------------+
| mambaout_small                     |                  |
+------------------------------------+------------------+
| mambaout_small_rw                  |                  |
+------------------------------------+------------------+
| mambaout_tiny                      |                  |
+------------------------------------+------------------+
| mvitv2_base                        |                  |
+------------------------------------+------------------+
| mvitv2_base_cls                    |                  |
+------------------------------------+------------------+
| mvitv2_huge_cls                    |                  |
+------------------------------------+------------------+
| mvitv2_large                       |                  |
+------------------------------------+------------------+
| mvitv2_large_cls                   |                  |
+------------------------------------+------------------+
| mvitv2_small                       |                  |
+------------------------------------+------------------+
| mvitv2_small_cls                   |                  |
+------------------------------------+------------------+
| mvitv2_tiny                        |                  |
+------------------------------------+------------------+
| nest_base                          |                  |
+------------------------------------+------------------+
| nest_base_jx                       |                  |
+------------------------------------+------------------+
| nest_small                         |                  |
+------------------------------------+------------------+
| nest_small_jx                      |                  |
+------------------------------------+------------------+
| nest_tiny                          |                  |
+------------------------------------+------------------+
| nest_tiny_jx                       |                  |
+------------------------------------+------------------+
| nextvit_base                       |                  |
+------------------------------------+------------------+
| nextvit_large                      |                  |
+------------------------------------+------------------+
| nextvit_small                      |                  |
+------------------------------------+------------------+
| poolformer_m36                     |                  |
+------------------------------------+------------------+
| poolformer_m48                     |                  |
+------------------------------------+------------------+
| poolformer_s12                     |                  |
+------------------------------------+------------------+
| poolformer_s24                     |                  |
+------------------------------------+------------------+
| poolformer_s36                     |                  |
+------------------------------------+------------------+
| poolformerv2_m36                   |                  |
+------------------------------------+------------------+
| poolformerv2_m48                   |                  |
+------------------------------------+------------------+
| poolformerv2_s12                   |                  |
+------------------------------------+------------------+
| poolformerv2_s24                   |                  |
+------------------------------------+------------------+
| poolformerv2_s36                   |                  |
+------------------------------------+------------------+
| pvt_v2_b0                          |                  |
+------------------------------------+------------------+
| pvt_v2_b1                          |                  |
+------------------------------------+------------------+
| pvt_v2_b2                          |                  |
+------------------------------------+------------------+
| pvt_v2_b2_li                       |                  |
+------------------------------------+------------------+
| pvt_v2_b3                          |                  |
+------------------------------------+------------------+
| pvt_v2_b4                          |                  |
+------------------------------------+------------------+
| pvt_v2_b5                          |                  |
+------------------------------------+------------------+
| rdnet_base                         |                  |
+------------------------------------+------------------+
| rdnet_large                        |                  |
+------------------------------------+------------------+
| rdnet_small                        |                  |
+------------------------------------+------------------+
| rdnet_tiny                         |                  |
+------------------------------------+------------------+
| repvit_m0_9                        |                  |
+------------------------------------+------------------+
| repvit_m1                          |                  |
+------------------------------------+------------------+
| repvit_m1_0                        |                  |
+------------------------------------+------------------+
| repvit_m1_1                        |                  |
+------------------------------------+------------------+
| repvit_m1_5                        |                  |
+------------------------------------+------------------+
| repvit_m2                          |                  |
+------------------------------------+------------------+
| repvit_m2_3                        |                  |
+------------------------------------+------------------+
| repvit_m3                          |                  |
+------------------------------------+------------------+
| sam2_hiera_base_plus               |                  |
+------------------------------------+------------------+
| sam2_hiera_large                   |                  |
+------------------------------------+------------------+
| sam2_hiera_small                   |                  |
+------------------------------------+------------------+
| sam2_hiera_tiny                    |                  |
+------------------------------------+------------------+
| swin_base_patch4_window7_224       |                  |
+------------------------------------+------------------+
| swin_base_patch4_window12_384      |                  |
+------------------------------------+------------------+
| swin_large_patch4_window7_224      |                  |
+------------------------------------+------------------+
| swin_large_patch4_window12_384     |                  |
+------------------------------------+------------------+
| swin_s3_base_224                   |                  |
+------------------------------------+------------------+
| swin_s3_small_224                  |                  |
+------------------------------------+------------------+
| swin_s3_tiny_224                   |                  |
+------------------------------------+------------------+
| swin_small_patch4_window7_224      |                  |
+------------------------------------+------------------+
| swin_tiny_patch4_window7_224       |                  |
+------------------------------------+------------------+
| swinv2_base_window8_256            |                  |
+------------------------------------+------------------+
| swinv2_base_window12_192           |                  |
+------------------------------------+------------------+
| swinv2_base_window12to16_192to256  |                  |
+------------------------------------+------------------+
| swinv2_base_window12to24_192to384  |                  |
+------------------------------------+------------------+
| swinv2_base_window16_256           |                  |
+------------------------------------+------------------+
| swinv2_cr_base_224                 |                  |
+------------------------------------+------------------+
| swinv2_cr_base_384                 |                  |
+------------------------------------+------------------+
| swinv2_cr_base_ns_224              |                  |
+------------------------------------+------------------+
| swinv2_cr_giant_224                |                  |
+------------------------------------+------------------+
| swinv2_cr_giant_384                |                  |
+------------------------------------+------------------+
| swinv2_cr_huge_224                 |                  |
+------------------------------------+------------------+
| swinv2_cr_huge_384                 |                  |
+------------------------------------+------------------+
| swinv2_cr_large_224                |                  |
+------------------------------------+------------------+
| swinv2_cr_large_384                |                  |
+------------------------------------+------------------+
| swinv2_cr_small_224                |                  |
+------------------------------------+------------------+
| swinv2_cr_small_384                |                  |
+------------------------------------+------------------+
| swinv2_cr_small_ns_224             |                  |
+------------------------------------+------------------+
| swinv2_cr_small_ns_256             |                  |
+------------------------------------+------------------+
| swinv2_cr_tiny_224                 |                  |
+------------------------------------+------------------+
| swinv2_cr_tiny_384                 |                  |
+------------------------------------+------------------+
| swinv2_cr_tiny_ns_224              |                  |
+------------------------------------+------------------+
| swinv2_large_window12_192          |                  |
+------------------------------------+------------------+
| swinv2_large_window12to16_192to256 |                  |
+------------------------------------+------------------+
| swinv2_large_window12to24_192to384 |                  |
+------------------------------------+------------------+
| swinv2_small_window8_256           |                  |
+------------------------------------+------------------+
| swinv2_small_window16_256          |                  |
+------------------------------------+------------------+
| swinv2_tiny_window8_256            |                  |
+------------------------------------+------------------+
| swinv2_tiny_window16_256           |                  |
+------------------------------------+------------------+
| tiny_vit_5m_224                    |                  |
+------------------------------------+------------------+
| tiny_vit_11m_224                   |                  |
+------------------------------------+------------------+
| tiny_vit_21m_224                   |                  |
+------------------------------------+------------------+
| tiny_vit_21m_384                   |                  |
+------------------------------------+------------------+
| tiny_vit_21m_512                   |                  |
+------------------------------------+------------------+
| tresnet_l                          |                  |
+------------------------------------+------------------+
| tresnet_m                          |                  |
+------------------------------------+------------------+
| tresnet_v2_l                       |                  |
+------------------------------------+------------------+
| tresnet_xl                         |                  |
+------------------------------------+------------------+
| twins_pcpvt_base                   |                  |
+------------------------------------+------------------+
| twins_pcpvt_large                  |                  |
+------------------------------------+------------------+
| twins_pcpvt_small                  |                  |
+------------------------------------+------------------+
| twins_svt_base                     |                  |
+------------------------------------+------------------+
| twins_svt_large                    |                  |
+------------------------------------+------------------+
| twins_svt_small                    |                  |
+------------------------------------+------------------+



================================================
FILE: docs/index.rst
================================================
.. Segmentation Models documentation master file, created by
   sphinx-quickstart on Fri Nov 27 00:00:20 2020.
   You can adapt this file completely to your liking, but it should at least
   contain the root `toctree` directive.

Welcome to Segmentation Models's documentation!
===============================================

.. toctree::
   :maxdepth: 2
   :caption: Contents:

   install
   quickstart
   models
   encoders
   encoders_timm
   losses
   metrics
   save_load
   insights


Indices and tables
==================

* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`


================================================
FILE: docs/insights.rst
================================================
💡 Insights
===========

1. Models architecture
~~~~~~~~~~~~~~~~~~~~~~

All segmentation models in SMP (this library short name) are made of:

- encoder (feature extractor, a.k.a backbone)
- decoder (features fusion block to create segmentation *mask*)
- segmentation head (final head to reduce number of channels from decoder and upsample mask to preserve input-output spatial resolution identity)
- classification head (optional head which build on top of deepest encoder features)


2. Creating your own encoder
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Encoder is a "classification model" which extract features from image and pass it to decoder.
Each encoder should have following attributes and methods and be inherited from `segmentation_models_pytorch.encoders._base.EncoderMixin`

.. code-block:: python

    class MyEncoder(torch.nn.Module, EncoderMixin):
        
        def __init__(self, **kwargs):
            super().__init__()
            
            # A number of channels for each encoder feature tensor, list of integers
            self._out_channels: List[int] = [3, 16, 64, 128, 256, 512]

            # A number of stages in decoder (in other words number of downsampling operations), integer
            # use in in forward pass to reduce number of returning features
            self._depth: int = 5 

            # Default number of input channels in first Conv2d layer for encoder (usually 3)
            self._in_channels: int = 3 
            
            # Define encoder modules below
            ...

        def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
            """Produce list of features of different spatial resolutions, each feature is a 4D torch.tensor of
            shape NCHW (features should be sorted in descending order according to spatial resolution, starting
            with resolution same as input `x` tensor).

            Input: `x` with shape (1, 3, 64, 64)
            Output: [f0, f1, f2, f3, f4, f5] - features with corresponding shapes
                    [(1, 3, 64, 64), (1, 64, 32, 32), (1, 128, 16, 16), (1, 256, 8, 8),
                    (1, 512, 4, 4), (1, 1024, 2, 2)] (C - dim may differ)

            also should support number of features according to specified depth, e.g. if depth = 5,
            number of feature tensors = 6 (one with same resolution as input and 5 downsampled),
            depth = 3 -> number of feature tensors = 4 (one with same resolution as input and 3 downsampled).
            """

            return [feat1, feat2, feat3, feat4, feat5, feat6]

When you write your own Encoder class register its build parameters

.. code-block:: python

    smp.encoders.encoders["my_awesome_encoder"] = {
        "encoder": MyEncoder, # encoder class here
        "pretrained_settings": {
            "imagenet": {
                "mean": [0.485, 0.456, 0.406],
                "std": [0.229, 0.224, 0.225],
                "url": "https://some-url.com/my-model-weights",
                "input_space": "RGB",
                "input_range": [0, 1],
            },
        },
        "params": {
            # init params for encoder if any
        },
    },

Now you can use your encoder

.. code-block:: python

    model = smp.Unet(encoder_name="my_awesome_encoder")

For better understanding see more examples of encoder in smp.encoders module.

.. note::

    If it works fine, don`t forget to contribute your work and make a PR to SMP 😉

3. Aux classification output
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

All models support ``aux_params`` parameter, which is default set to ``None``. 
If ``aux_params = None`` than classification auxiliary output is not created, else
model produce not only ``mask``, but also ``label`` output with shape ``(N, C)``.

Classification head consist of following layers:
    
1. GlobalPooling
2. Dropout (optional)
3. Linear
4. Activation (optional)

Example:

.. code-block:: python
    
    aux_params=dict(
        pooling='avg',             # one of 'avg', 'max'
        dropout=0.5,               # dropout ratio, default is None
        activation='sigmoid',      # activation function, default is None
        classes=4,                 # define number of output labels
    )

    model = smp.Unet('resnet34', classes=4, aux_params=aux_params)
    mask, label = model(x)

    mask.shape, label.shape
    # (N, 4, H, W), (N, 4)

4. Freezing and unfreezing the encoder
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Sometimes you may want to freeze the encoder during training, e.g. when using pretrained backbones and only fine-tuning the decoder and segmentation head. 

All segmentation models in SMP provide two helper methods:

.. code-block:: python

    model = smp.Unet("resnet34", classes=2)

    # Freeze encoder: stops gradient updates and freezes normalization layer stats
    model.freeze_encoder()

    # Unfreeze encoder: re-enables training for encoder parameters and normalization layers
    model.unfreeze_encoder()

.. important::
    - Freezing sets ``requires_grad = False`` for all encoder parameters.
    - Normalization layers that track running statistics (e.g., BatchNorm and InstanceNorm layers) are set to ``.eval()`` mode to prevent updates to ``running_mean`` and ``running_var``.
    - If you later call ``model.train()``, frozen encoders will remain frozen until you call ``unfreeze_encoder()``.


================================================
FILE: docs/install.rst
================================================
⚙️ Installation
===============

PyPI version:

.. code-block:: bash

    $ pip install -U segmentation-models-pytorch


Latest version from source:

.. code-block:: bash

    $ pip install -U git+https://github.com/qubvel/segmentation_models.pytorch

================================================
FILE: docs/losses.rst
================================================
📉 Losses
=========

Collection of popular semantic segmentation losses. Adapted from
an awesome repo with pytorch utils https://github.com/BloodAxe/pytorch-toolbelt

Constants
~~~~~~~~~
.. automodule:: segmentation_models_pytorch.losses.constants
        :members:

JaccardLoss
~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.JaccardLoss

DiceLoss
~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.DiceLoss

TverskyLoss
~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.TverskyLoss

FocalLoss
~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.FocalLoss

LovaszLoss
~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.LovaszLoss

SoftBCEWithLogitsLoss
~~~~~~~~~~~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.SoftBCEWithLogitsLoss

SoftCrossEntropyLoss
~~~~~~~~~~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.SoftCrossEntropyLoss

MCCLoss
~~~~~~~~~~~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.losses.MCCLoss
        :members: forward


================================================
FILE: docs/make.bat
================================================
@ECHO OFF

pushd %~dp0

REM Command file for Sphinx documentation

if "%SPHINXBUILD%" == "" (
	set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=source
set BUILDDIR=build

if "%1" == "" goto help

%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
	echo.
	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
	echo.installed, then set the SPHINXBUILD environment variable to point
	echo.to the full path of the 'sphinx-build' executable. Alternatively you
	echo.may add the Sphinx directory to PATH.
	echo.
	echo.If you don't have Sphinx installed, grab it from
	echo.http://sphinx-doc.org/
	exit /b 1
)

%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end

:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%

:end
popd


================================================
FILE: docs/metrics.rst
================================================
📏 Metrics
==========

Functional metrics
~~~~~~~~~~~~~~~~~~

.. currentmodule:: segmentation_models_pytorch.metrics.functional

.. autosummary::

   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

.. automodule:: segmentation_models_pytorch.metrics.functional
   :members:


================================================
FILE: docs/models.rst
================================================
🕸️ Segmentation Models
==============================


.. contents::
   :local:

.. _unet:

Unet
~~~~
.. autoclass:: segmentation_models_pytorch.Unet


.. _unetplusplus:

Unet++
~~~~~~
.. autoclass:: segmentation_models_pytorch.UnetPlusPlus


.. _fpn:

FPN
~~~
.. autoclass:: segmentation_models_pytorch.FPN


.. _pspnet:

PSPNet
~~~~~~
.. autoclass:: segmentation_models_pytorch.PSPNet


.. _deeplabv3:

DeepLabV3
~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.DeepLabV3


.. _deeplabv3plus:

DeepLabV3+
~~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.DeepLabV3Plus


.. _linknet:

Linknet
~~~~~~~
.. autoclass:: segmentation_models_pytorch.Linknet


.. _manet:

MAnet
~~~~~~
.. autoclass:: segmentation_models_pytorch.MAnet


.. _pan:

PAN
~~~
.. autoclass:: segmentation_models_pytorch.PAN


.. _upernet:

UPerNet
~~~~~~~
.. autoclass:: segmentation_models_pytorch.UPerNet


.. _segformer:

Segformer
~~~~~~~~~
.. autoclass:: segmentation_models_pytorch.Segformer


.. _dpt:

DPT
~~~

.. note::

    See full list of DPT-compatible timm encoders in :ref:`dpt-encoders`.

.. note::

    For some encoders, the model requires ``dynamic_img_size=True`` to be passed in order to work with resolutions different from what the encoder was trained for.

.. autoclass:: segmentation_models_pytorch.DPT


================================================
FILE: docs/quickstart.rst
================================================
🚀 Quick Start
==============

**1. Create segmentation model**

Segmentation model is just a PyTorch nn.Module, which can be created as easy as:

.. code-block:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet(
        encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
        encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
        in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
        classes=3,                      # model output channels (number of classes in your dataset)
    )

- Check the page with available :doc:`model architectures <models>`.
- Check the table with :doc:`available ported encoders and its corresponding weights <encoders>`.
- `Pytorch Image Models (timm) <https://github.com/huggingface/pytorch-image-models>`_ encoders are also supported, check it :doc:`here<encoders_timm>`.

Alternatively, you can use `smp.create_model` function to create a model by name:

.. code-block:: python

    model = smp.create_model(
        arch="fpn",                     # name of the architecture, e.g. 'Unet'/ 'FPN' / etc. Case INsensitive!
        encoder_name="mit_b0",
        encoder_weights="imagenet",
        in_channels=1,
        classes=3,
    )


**2. Configure data preprocessing**

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give your better results (higher metric score and faster convergence). But it is relevant only for 1-2-3-channels images and **not necessary** in case you train the whole model, not only decoder.

.. code-block:: python

    from segmentation_models_pytorch.encoders import get_preprocessing_fn

    preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')


**3. Congratulations!** 🎉


You are done! Now you can train your model with your favorite framework, or as simple as:

.. code-block:: python

    for images, gt_masks in dataloader:

        predicted_mask = model(images)
        loss = loss_fn(predicted_mask, gt_masks)

        loss.backward()
        optimizer.step()

Check the following examples:

.. |colab-badge| image:: https://colab.research.google.com/assets/colab-badge.svg
   :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
   :alt: Open In Colab

- Finetuning notebook on Oxford Pet dataset with `PyTorch Lightning <https://github.com/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb>`__ |colab-badge|
- Finetuning script for cloth segmentation with `PyTorch Lightning <https://github.com/ternaus/cloths_segmentation>`__


================================================
FILE: docs/save_load.rst
================================================
📂 Saving and Loading
=====================

In this section, we will discuss how to save a trained model, push it to the Hugging Face Hub, and load it back for later use.

Saving and Sharing a Model
--------------------------

Once you have trained your model, you can save it using the `.save_pretrained` method. This method saves the model configuration and weights to a directory of your choice.
And, optionally, you can push the model to the Hugging Face Hub by setting the `push_to_hub` parameter to `True`.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet('resnet34', encoder_weights='imagenet')

    # After training your model, save it to a directory
    model.save_pretrained('./my_model')

    # Or saved and pushed to the Hub simultaneously
    model.save_pretrained('username/my-model', push_to_hub=True)

Loading Trained Model
---------------------

Once your model is saved and pushed to the Hub, you can load it back using the `smp.from_pretrained` method. This method allows you to load the model weights and configuration from a directory or directly from the Hub.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    # Load the model from the local directory
    model = smp.from_pretrained('./my_model')

    # Alternatively, load the model directly from the Hugging Face Hub
    model = smp.from_pretrained('username/my-model')

Loading pre-trained model with different number of classes for fine-tuning:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.from_pretrained('<path-or-repo-name>', classes=5, strict=False)

Saving model Metrics and Dataset Name
-------------------------------------

You can simply pass the `metrics` and `dataset` parameters to the `save_pretrained` method to save the model metrics and dataset name in Model Card along with the model configuration and weights.

For example:

.. code:: python

    import segmentation_models_pytorch as smp

    model = smp.Unet('resnet34', encoder_weights='imagenet')

    # After training your model, save it to a directory
    model.save_pretrained('./my_model', metrics={'accuracy': 0.95}, dataset='my_dataset')

    # Or saved and pushed to the Hub simultaneously
    model.save_pretrained('username/my-model', push_to_hub=True, metrics={'accuracy': 0.95}, dataset='my_dataset')

Saving with preprocessing transform (Albumentations)
----------------------------------------------------

You can save the preprocessing transform along with the model and push it to the Hub.
This can be useful when you want to share the model with the preprocessing transform that was used during training,
to make sure that the inference pipeline is consistent with the training pipeline.

.. code:: python

    import albumentations as A
    import segmentation_models_pytorch as smp

    # Define a preprocessing transform for image that would be used during inference
    preprocessing_transform = A.Compose([
        A.Resize(256, 256),
        A.Normalize()
    ])

    model = smp.Unet()

    directory_or_repo_on_the_hub = "qubvel-hf/unet-with-transform"  # <username>/<repo-name>

    # Save the model and transform (and pus ot hub, if needed)
    model.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)
    preprocessing_transform.save_pretrained(directory_or_repo_on_the_hub, push_to_hub=True)

    # Loading transform and model
    restored_model = smp.from_pretrained(directory_or_repo_on_the_hub)
    restored_transform = A.Compose.from_pretrained(directory_or_repo_on_the_hub)

    print(restored_transform)

Conclusion
----------

By following these steps, you can easily save, share, and load your models, facilitating collaboration and reproducibility in your projects. Don't forget to replace the placeholders with your actual model paths and names.

|binary-segmentation-intro|
|save-load-model-and-share-with-hf-hub|

.. |binary-segmentation-intro| image:: https://colab.research.google.com/assets/colab-badge.svg
    :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb
    :alt: Open In Colab

.. |save-load-model-and-share-with-hf-hub| image:: https://colab.research.google.com/assets/colab-badge.svg
    :target: https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/save_load_model_and_share_with_hf_hub.ipynb
    :alt: Open In Colab


================================================
FILE: examples/binary_segmentation_buildings.py
================================================
"""
This script demonstrates how to train a binary segmentation model using the
CamVid dataset and segmentation_models_pytorch. The CamVid dataset is a
collection of videos with pixel-level annotations for semantic segmentation.
The dataset includes 367 training images, 101 validation images, and 233 test.
Each training image has a corresponding mask that labels each pixel as belonging
to these classes with the numerical labels as follows:
- Sky: 0
- Building: 1
- Pole: 2
- Road: 3
- Pavement: 4
- Tree: 5
- SignSymbol: 6
- Fence: 7
- Car: 8
- Pedestrian: 9
- Bicyclist: 10
- Unlabelled: 11

In this script, we focus on binary segmentation, where the goal is to classify
each pixel as whether belonging to a certain class (Foregorund) or
not (Background).

Class Labels:
- 0: Background
- 1: Foreground

The script includes the following steps:
1. Set the device to GPU if available, otherwise use CPU.
2. Download the CamVid dataset if it is not already present.
3. Define hyperparameters for training.
4. Define a custom dataset class for loading and preprocessing the CamVid
     dataset.
5. Define a function to visualize images and masks.
6. Create datasets and dataloaders for training, validation, and testing.
7. Define a model class for the segmentation task.
8. Train the model using the training and validation datasets.
9. Evaluate the model using the test dataset and save the output masks and
     metrics.
"""

import logging
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset
from tqdm import tqdm

import segmentation_models_pytorch as smp

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(message)s",
    datefmt="%d:%m:%Y %H:%M:%S",
)

# ----------------------------
# Set the device to GPU if available
# ----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
logging.info(f"Using device: {device}")
if device == "cpu":
    os.system("export OMP_NUM_THREADS=64")
    torch.set_num_threads(os.cpu_count())

# ----------------------------
# Download the CamVid dataset, if needed
# ----------------------------
# Change this to your desired directory
main_dir = "./examples/binary_segmentation_data/"

data_dir = os.path.join(main_dir, "dataset")
if not os.path.exists(data_dir):
    logging.info("Loading data...")
    os.system(f"git clone https://github.com/alexgkendall/SegNet-Tutorial {data_dir}")
    logging.info("Done!")

# Create a directory to store the output masks
output_dir = os.path.join(main_dir, "output_images")
os.makedirs(output_dir, exist_ok=True)

# ----------------------------
# Define the hyperparameters
# ----------------------------
epochs_max = 200  # Number of epochs to train the model
adam_lr = 2e-4  # Learning rate for the Adam optimizer
eta_min = 1e-5  # Minimum learning rate for the scheduler
batch_size = 8  # Batch size for training
input_image_reshape = (320, 320)  # Desired shape for the input images and masks
foreground_class = 1  # 1 for binary segmentation


# ----------------------------
# Define a custom dataset class for the CamVid dataset
# ----------------------------
class Dataset(BaseDataset):
    """
    A custom dataset class for binary segmentation tasks.

    Parameters:
    ----------

    - images_dir (str): Directory containing the input images.
    - masks_dir (str): Directory containing the corresponding masks.
    - input_image_reshape (tuple, optional): Desired shape for the input
      images and masks. Default is (320, 320).
    - foreground_class (int, optional): The class value in the mask to be
      considered as the foreground. Default is 1.
    - augmentation (callable, optional): A function/transform to apply to the
      images and masks for data augmentation.
    """

    def __init__(
        self,
        images_dir,
        masks_dir,
        input_image_reshape=(320, 320),
        foreground_class=1,
        augmentation=None,
    ):
        self.ids = os.listdir(images_dir)
        self.images_filepaths = [
            os.path.join(images_dir, image_id) for image_id in self.ids
        ]
        self.masks_filepaths = [
            os.path.join(masks_dir, image_id) for image_id in self.ids
        ]

        self.input_image_reshape = input_image_reshape
        self.foreground_class = foreground_class
        self.augmentation = augmentation

    def __getitem__(self, i):
        """
        Retrieves the image and corresponding mask at index `i`.

        Parameters:
        ----------

        - i (int): Index of the image and mask to retrieve.
        Returns:
        - A tuple containing:
            - image (torch.Tensor): The preprocessed image tensor of shape
            (1, input_image_reshape) - e.g., (1, 320, 320) - normalized to [0, 1].
            - mask_remap (torch.Tensor): The preprocessed mask tensor of
            shape input_image_reshape with values 0 or 1.
        """
        # Read the image
        image = cv2.imread(
            self.images_filepaths[i], cv2.IMREAD_GRAYSCALE
        )  # Read image as grayscale
        image = np.expand_dims(image, axis=-1)  # Add channel dimension

        # resize image to input_image_reshape
        image = cv2.resize(image, self.input_image_reshape)

        # Read the mask in grayscale mode
        mask = cv2.imread(self.masks_filepaths[i], 0)

        # Update the mask: Set foreground_class to 1 and the rest to 0
        mask_remap = np.where(mask == self.foreground_class, 1, 0).astype(np.uint8)

        # resize mask to input_image_reshape
        mask_remap = cv2.resize(mask_remap, self.input_image_reshape)

        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask_remap)
            image, mask_remap = sample["image"], sample["mask"]

        # Convert to PyTorch tensors
        # Add channel dimension if missing
        if image.ndim == 2:
            image = np.expand_dims(image, axis=-1)

        # HWC -> CHW and normalize to [0, 1]
        image = torch.tensor(image).float().permute(2, 0, 1) / 255.0

        # Ensure mask is LongTensor
        mask_remap = torch.tensor(mask_remap).long()

        return image, mask_remap

    def __len__(self):
        return len(self.ids)


# Define a class for the CamVid model
class CamVidModel(torch.nn.Module):
    """
    A PyTorch model for binary segmentation using the Segmentation Models
    PyTorch library.

    Parameters:
    ----------

    - arch (str): The architecture name of the segmentation model
       (e.g., 'Unet', 'FPN').
    - encoder_name (str): The name of the encoder to use
       (e.g., 'resnet34', 'vgg16').
    - in_channels (int, optional): Number of input channels (e.g., 3 for RGB).
    - out_classes (int, optional): Number of output classes (e.g., 1 for binary)
    **kwargs: Additional keyword arguments to pass to the model
    creation function.
    """

    def __init__(self, arch, encoder_name, in_channels=3, out_classes=1, **kwargs):
        super().__init__()
        self.mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
        self.std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
        self.model = smp.create_model(
            arch,
            encoder_name=encoder_name,
            in_channels=in_channels,
            classes=out_classes,
            **kwargs,
        )

    def forward(self, image):
        # Normalize image
        image = (image - self.mean) / self.std
        mask = self.model(image)
        return mask


def visualize(output_dir, image_filename, **images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(" ".join(name.split("_")).title())
        plt.imshow(image)
    plt.show()
    plt.savefig(os.path.join(output_dir, image_filename))
    plt.close()


# Use multiple CPUs in parallel
def train_and_evaluate_one_epoch(
    model, train_dataloader, valid_dataloader, optimizer, scheduler, loss_fn, device
):
    # Set the model to training mode
    model.train()
    train_loss = 0
    for batch in tqdm(train_dataloader, desc="Training"):
        images, masks = batch
        images, masks = images.to(device), masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)

        loss = loss_fn(outputs, masks)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    scheduler.step()
    avg_train_loss = train_loss / len(train_dataloader)

    # Set the model to evaluation mode
    model.eval()
    val_loss = 0
    with torch.inference_mode():
        for batch in tqdm(valid_dataloader, desc="Evaluating"):
            images, masks = batch
            images, masks = images.to(device), masks.to(device)

            outputs = model(images)
            loss = loss_fn(outputs, masks)

            val_loss += loss.item()

    avg_val_loss = val_loss / len(valid_dataloader)
    return avg_train_loss, avg_val_loss


def train_model(
    model,
    train_dataloader,
    valid_dataloader,
    optimizer,
    scheduler,
    loss_fn,
    device,
    epochs,
):
    train_losses = []
    val_losses = []

    for epoch in range(epochs):
        avg_train_loss, avg_val_loss = train_and_evaluate_one_epoch(
            model,
            train_dataloader,
            valid_dataloader,
            optimizer,
            scheduler,
            loss_fn,
            device,
        )
        train_losses.append(avg_train_loss)
        val_losses.append(avg_val_loss)

        logging.info(
            f"Epoch {epoch + 1}/{epochs}, Training Loss: {avg_train_loss:.2f}, Validation Loss: {avg_val_loss:.2f}"
        )

    history = {
        "train_losses": train_losses,
        "val_losses": val_losses,
    }
    return history


def test_model(model, output_dir, test_dataloader, loss_fn, device):
    # Set the model to evaluation mode
    model.eval()
    test_loss = 0
    tp, fp, fn, tn = 0, 0, 0, 0
    with torch.inference_mode():
        for batch in tqdm(test_dataloader, desc="Evaluating"):
            images, masks = batch
            images, masks = images.to(device), masks.to(device)

            outputs = model(images)
            loss = loss_fn(outputs, masks)

            for i, output in enumerate(outputs):
                input = images[i].cpu().numpy().transpose(1, 2, 0)
                output = output.squeeze().cpu().numpy()

                visualize(
                    output_dir,
                    f"output_{i}.png",
                    input_image=input,
                    output_mask=output,
                    binary_mask=output > 0.5,
                )

            test_loss += loss.item()

            prob_mask = outputs.sigmoid().squeeze(1)
            pred_mask = (prob_mask > 0.5).long()
            batch_tp, batch_fp, batch_fn, batch_tn = smp.metrics.get_stats(
                pred_mask, masks, mode="binary"
            )
            tp += batch_tp.sum().item()
            fp += batch_fp.sum().item()
            fn += batch_fn.sum().item()
            tn += batch_tn.sum().item()

        test_loss_mean = test_loss / len(test_dataloader)
        logging.info(f"Test Loss: {test_loss_mean:.2f}")

    iou_score = smp.metrics.iou_score(
        torch.tensor([tp]),
        torch.tensor([fp]),
        torch.tensor([fn]),
        torch.tensor([tn]),
        reduction="micro",
    )

    return test_loss_mean, iou_score.item()


# ----------------------------
# Define the data directories and create the datasets
# ----------------------------
x_train_dir = os.path.join(data_dir, "CamVid", "train")
y_train_dir = os.path.join(data_dir, "CamVid", "trainannot")

x_val_dir = os.path.join(data_dir, "CamVid", "val")
y_val_dir = os.path.join(data_dir, "CamVid", "valannot")

x_test_dir = os.path.join(data_dir, "CamVid", "test")
y_test_dir = os.path.join(data_dir, "CamVid", "testannot")

train_dataset = Dataset(
    x_train_dir,
    y_train_dir,
    input_image_reshape=input_image_reshape,
    foreground_class=foreground_class,
)
valid_dataset = Dataset(
    x_val_dir,
    y_val_dir,
    input_image_reshape=input_image_reshape,
    foreground_class=foreground_class,
)
test_dataset = Dataset(
    x_test_dir,
    y_test_dir,
    input_image_reshape=input_image_reshape,
    foreground_class=foreground_class,
)

image, mask = train_dataset[0]
logging.info(f"Unique values in mask: {np.unique(mask)}")
logging.info(f"Image shape: {image.shape}")
logging.info(f"Mask shape: {mask.shape}")

# ----------------------------
# Create the dataloaders using the datasets
# ----------------------------
logging.info(f"Train size: {len(train_dataset)}")
logging.info(f"Valid size: {len(valid_dataset)}")
logging.info(f"Test size: {len(test_dataset)}")

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# ----------------------------
# Lets look at some samples
# ----------------------------
# Visualize and save train sample
sample = train_dataset[0]
visualize(
    output_dir,
    "train_sample.png",
    train_image=sample[0].numpy().transpose(1, 2, 0),
    train_mask=sample[1].squeeze(),
)

# Visualize and save validation sample
sample = valid_dataset[0]
visualize(
    output_dir,
    "validation_sample.png",
    validation_image=sample[0].numpy().transpose(1, 2, 0),
    validation_mask=sample[1].squeeze(),
)

# Visualize and save test sample
sample = test_dataset[0]
visualize(
    output_dir,
    "test_sample.png",
    test_image=sample[0].numpy().transpose(1, 2, 0),
    test_mask=sample[1].squeeze(),
)

# ----------------------------
# Create and train the model
# ----------------------------
max_iter = epochs_max * len(train_dataloader)  # Total number of iterations

model = CamVidModel("Unet", "resnet34", in_channels=3, out_classes=1)

# Training loop
model = model.to(device)

# Define the Adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=adam_lr)

# Define the learning rate scheduler
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_iter, eta_min=eta_min)

# Define the loss function
loss_fn = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True)

# Train the model
history = train_model(
    model,
    train_dataloader,
    valid_dataloader,
    optimizer,
    scheduler,
    loss_fn,
    device,
    epochs_max,
)

# Visualize the training and validation losses
plt.figure(figsize=(10, 5))
plt.plot(history["train_losses"], label="Train Loss")
plt.plot(history["val_losses"], label="Validation Loss")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.title("Training and Validation Losses")
plt.legend()
plt.savefig(os.path.join(output_dir, "train_val_losses.png"))
plt.close()


# Evaluate the model
test_loss = test_model(model, output_dir, test_dataloader, loss_fn, device)

logging.info(f"Test Loss: {test_loss[0]}, IoU Score: {test_loss[1]}")
logging.info(f"The output masks are saved in {output_dir}.")


================================================
FILE: examples/binary_segmentation_intro.ipynb
================================================
{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U3WUb8t2P2e5"
      },
      "source": [
        "\ud83c\udded \ud83c\uddea \ud83c\uddf1 \ud83c\uddf1 \ud83c\uddf4 \ud83d\udc4b\n",
        "\n",
        "This example shows how to use `segmentation-models-pytorch` for **binary** semantic segmentation. We will use the [The Oxford-IIIT Pet Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/) (this is an adopted example from Albumentations package [docs](https://albumentations.ai/docs/examples/pytorch_semantic_segmentation/), which is strongly recommended to read, especially if you never used this package for augmentations before). \n",
        "\n",
        "The task will be to classify each pixel of an input image either as pet \ud83d\udc36\ud83d\udc31 or as a background.\n",
        "\n",
        "\n",
        "What we are going to overview in this example:  \n",
        "\n",
        " - \ud83d\udcdc `Datasets` and `DataLoaders` preparation (with predefined dataset class).  \n",
        " - \ud83d\udce6 `LightningModule` preparation: defining training, validation and test routines.  \n",
        " - \ud83d\udcc8 Writing `IoU` metric inside the `LightningModule` for measuring quality of segmentation.  \n",
        " - \ud83d\udc36 Results visualization.\n",
        "\n",
        "\n",
        "> It is expected you are familiar with Python, PyTorch and have some experience with training neural networks before!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2024-08-18T04:37:36.751747Z",
          "iopub.status.busy": "2024-08-18T04:37:36.750812Z",
          "iopub.status.idle": "2024-08-18T04:38:26.758872Z",
          "shell.execute_reply": "2024-08-18T04:38:26.757586Z",
          "shell.execute_reply.started": "2024-08-18T04:37:36.751710Z"
        },
        "id": "DYNdz8s56qOu",
        "outputId": "7f343747-532d-417c-fc72-fda5c713d4e3",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install -U git+https://github.com/qubvel-org/segmentation_models.pytorch\n",
        "!pip install lightning albumentations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2024-08-18T04:38:26.761388Z",
          "iopub.status.busy": "2024-08-18T04:38:26.761047Z",
          "iopub.s