Enhancements

.gitignore

@@ -0,0 +1,168 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-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/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# custom
+data/
+deploy/
+nuitka-crash-report.xml
+outputs/
+pretrained/
+

README.md

@@ -196,11 +196,11 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py
 
 ## Knowledge distillation
 
-| Model | Encoder       | Decoder                 | kd_training | mIoU(200 epoch) | mIoU(800 epoch) |
-|:-----:|:-------------:|:-----------------------:|:-----------:|:---------------:|:---------------:|
-| SMP   | DeepLabv3Plus | ResNet-101 <br> teacher | -           | 78.10           | 79.20           |
-| SMP   | DeepLabv3Plus | ResNet-18 <br> student  | False       | 73.97           | 75.90           |
-| SMP   | DeepLabv3Plus | ResNet-18 <br> student  | True        | 75.20           | 76.41           |
+| Model | Encoder                 | Decoder       | kd_training | mIoU(200 epoch) | mIoU(800 epoch) |
+|:-----:|:-----------------------:|:-------------:|:-----------:|:---------------:|:---------------:|
+| SMP   | ResNet-101 <br> teacher | DeepLabv3Plus | -           | 78.10           | 79.20           |
+| SMP   | ResNet-18 <br> student  | DeepLabv3Plus | False       | 73.97           | 75.90           |
+| SMP   | ResNet-18 <br> student  | DeepLabv3Plus | True        | 75.20           | 76.41           |
 
 # Prepare the dataset
 

configs/__init__.py

@@ -1,2 +1,3 @@
+from .base_config import BaseConfig
 from .my_config import MyConfig
 from .parser import load_parser

configs/parser.py

@@ -16,7 +16,7 @@ def load_parser(config):
 def get_parser():
     parser = argparse.ArgumentParser()
     # Dataset
-    parser.add_argument('--dataset', type=str, default=None, choices=['cityscapes'],
+    parser.add_argument('--dataset', type=str, default=None, choices=['cityscapes', 'custom'],
         help='choose which dataset you want to use')
     parser.add_argument('--dataroot', type=str, default=None, 
         help='path to your dataset')

core/seg_trainer.py

@@ -158,6 +158,10 @@ def predict(self, config):
 
         self.logger.info('\nStart predicting...\n')
 
+        save_images_dir = os.path.join(config.save_dir, 'predicts')
+        if not os.path.exists(save_images_dir):
+            os.makedirs(save_images_dir)
+
         self.model.eval() # Put model in evalation mode
 
         for (images, images_aug, img_names) in tqdm(self.test_loader):
@@ -171,7 +175,7 @@ def predict(self, config):
 
             # Saving results
             for i in range(preds.shape[0]):
-                save_path = os.path.join(config.save_dir, img_names[i])
+                save_path = os.path.join(save_images_dir, img_names[i])
                 save_suffix = img_names[i].split('.')[-1]
 
                 pred = Image.fromarray(preds[i].astype(np.uint8))

datasets/__init__.py

@@ -1,7 +1,11 @@
 from torch.utils.data import DataLoader
 from .cityscapes import Cityscapes
+from .custom import Custom
 
-dataset_hub = {'cityscapes':Cityscapes,}
+dataset_hub = {
+    'cityscapes':Cityscapes,
+    'custom': Custom,
+}
 
 
 def get_dataset(config):

datasets/custom.py

@@ -0,0 +1,85 @@
+import os
+from collections import namedtuple
+import yaml
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+import albumentations as AT
+from albumentations.pytorch import ToTensorV2
+from utils import transforms
+
+
+class Custom(Dataset):
+    '''
+    demo for load custom datasets.
+    '''
+
+    def __init__(self, config, mode='train'):
+        data_root = os.path.expanduser(config.data_root)
+        dataset_config_filepath = os.path.join(data_root, 'data.yaml')
+        if not os.path.exists(dataset_config_filepath):
+            raise Exception(f"{dataset_config_filepath} not exists.")
+        with open(dataset_config_filepath, 'r', encoding='utf-8') as yaml_file:
+            dataset_config = yaml.safe_load(yaml_file)
+        print('dataset_config: ', dataset_config)
+        data_root = dataset_config['path']
+        # self.num_classes = len(dataset_config['names'])
+        self.id_to_train_id = dict()
+        for i in range(len(dataset_config['names'])):
+            self.id_to_train_id[i] = i
+        # self.train_id_to_name = dict()
+        # for k, v in dataset_config['names'].items():
+        #     self.train_id_to_name[k] = str(v)
+
+        img_dir = os.path.join(data_root, mode, 'imgs')
+        msk_dir = os.path.join(data_root, mode, 'masks')
+
+        if not os.path.isdir(img_dir):
+            raise RuntimeError(f'Image directory: {img_dir} does not exist.')
+
+        if not os.path.isdir(msk_dir):
+            raise RuntimeError(f'Mask directory: {msk_dir} does not exist.')
+
+        if mode == 'train':
+            self.transform = AT.Compose([
+                transforms.ResizeToSquare(size=config.train_size),
+                transforms.Scale(scale=config.scale),
+                AT.RandomScale(scale_limit=config.randscale),
+                AT.PadIfNeeded(min_height=config.crop_h, min_width=config.crop_w, value=(114,114,114), mask_value=(0,0,0)),
+                AT.RandomCrop(height=config.crop_h, width=config.crop_w),
+                AT.ColorJitter(brightness=config.brightness, contrast=config.contrast, saturation=config.saturation),
+                AT.HorizontalFlip(p=config.h_flip),
+                AT.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0)),
+                ToTensorV2(),                
+            ])
+
+        elif mode == 'val':
+            self.transform = AT.Compose([
+                transforms.ResizeToSquare(size=config.test_size),
+                transforms.Scale(scale=config.scale),
+                AT.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0)),
+                ToTensorV2(),
+            ])
+
+        self.images = []
+        self.masks = []
+
+        for img_file_name in os.listdir(img_dir):
+            img_file_basename = os.path.splitext(img_file_name)[0]
+
+            self.images.append(os.path.join(img_dir, img_file_name))
+            self.masks.append(os.path.join(msk_dir, img_file_basename + '.png'))
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        image = np.asarray(Image.open(self.images[index]).convert('RGB'))
+        mask = np.asarray(Image.open(self.masks[index]).convert('L'))
+
+        # Perform augmentation and normalization
+        augmented = self.transform(image=image, mask=mask)
+        image, mask = augmented['image'], augmented['mask']
+
+        return image, mask
+

models/__init__.py

@@ -68,7 +68,15 @@ def get_model(config):
         if config.decoder not in decoder_hub:
             raise ValueError(f"Unsupported decoder type: {config.decoder}")
 
-        model = decoder_hub[config.decoder](encoder_name=config.encoder, 
+        if config.encoder.startswith('mit_b') and config.decoder in ['pan']:
+                model = decoder_hub[config.decoder](encoder_name=config.encoder, 
+                                                encoder_weights=config.encoder_weights, 
+                                                encoder_output_stride=32,
+                                                in_channels=3, classes=config.num_class)
+        elif config.encoder.startswith('mit_b') and config.decoder in ['deeplabv3', 'deeplabv3p', 'linknet', 'unetpp']:
+            raise ValueError("Encoder `{}` is not supported for `{}".format(config.encoder, config.decoder))
+        else:
+            model = decoder_hub[config.decoder](encoder_name=config.encoder, 
                                             encoder_weights=config.encoder_weights, 
                                             in_channels=3, classes=config.num_class)
 

models/ddrnet.py

@@ -52,6 +52,11 @@ def forward(self, x, is_training=False):
         x = self.seg_head(x)
         x = F.interpolate(x, size, mode='bilinear', align_corners=True)
 
+        if torch.onnx.is_in_onnx_export():
+            # output_data = x.softmax(dim=1)
+            max_probs, predictions = x.max(1, keepdim=True)
+            return predictions.to(torch.int8)
+
         if self.use_aux and is_training:
             return x, (x_aux,)
         else:

models/stdc.py

@@ -87,6 +87,11 @@ def forward(self, x, is_training=False):
         x = self.seg_head(x)
         x = F.interpolate(x, size, mode='bilinear', align_corners=True)
 
+        if torch.onnx.is_in_onnx_export():
+            # output_data = x.softmax(dim=1)
+            max_probs, predictions = x.max(1, keepdim=True)
+            return predictions.to(torch.int8)
+
         if self.use_detail_head and is_training:
             x_detail = self.detail_head(x3)
             return x, x_detail

requirements.txt

@@ -0,0 +1,11 @@
+torch # == 1.8.1
+segmentation-models-pytorch
+torchmetrics
+albumentations
+loguru
+tqdm
+tensorboard
+onnx
+onnxsim
+onnxruntime
+labelme