Merge pull request #7 from amitklinger/Hailo-1.0

Hailo 1.0

configs/_base_/datasets/cityscapes10classes.py

@@ -0,0 +1,67 @@
+# dataset settings
+dataset_type = 'Cityscapes10classesDataset'
+data_root = '/data/data/cityscapes10classes/'
+crop_size = (512, 1024)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(
+        type='RandomResize',
+        scale=(2048,1024),
+        ratio_range=(0.5, 2.0),
+        keep_ratio=True),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', prob=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='PackSegInputs')
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='Resize', scale=(960, 736), keep_ratio=True),
+    # add loading annotation after ``Resize`` because ground truth
+    # does not need to do resize data transform
+    dict(type='LoadAnnotations'),
+    dict(type='PackSegInputs')
+]
+img_ratios = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
+tta_pipeline = [
+    dict(type='LoadImageFromFile', backend_args=None),
+    dict(
+        type='TestTimeAug',
+        transforms=[
+            [
+                dict(type='Resize', scale_factor=r, keep_ratio=True)
+                for r in img_ratios
+            ],
+            [
+                dict(type='RandomFlip', prob=0., direction='horizontal'),
+                dict(type='RandomFlip', prob=1., direction='horizontal')
+            ], [dict(type='LoadAnnotations')], [dict(type='PackSegInputs')]
+        ])
+]
+train_dataloader = dict(
+    batch_size=2,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        data_root=data_root,
+        data_prefix=dict(
+            img_path='leftImg8bit/train', seg_map_path='gtFine/train'),
+        pipeline=train_pipeline))
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        data_root=data_root,
+        data_prefix=dict(
+            img_path='leftImg8bit/val', seg_map_path='gtFine/val'),
+        pipeline=test_pipeline))
+test_dataloader = val_dataloader
+
+val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
+test_evaluator = val_evaluator

configs/fcn/fcn_hailo_10_classes_pp.py

@@ -0,0 +1,75 @@
+# model settings
+_base_ = [
+    '../_base_/datasets/cityscapes10classes.py', '../_base_/default_runtime.py',
+]
+
+# optimizer
+optimizer = dict(type='Adam', lr=0.001, weight_decay=1e-5)
+optim_wrapper = dict(type='OptimWrapper', optimizer=optimizer, clip_grad=None)
+
+# learning policy
+param_scheduler = [
+	dict(
+		type='LinearLR', start_factor=0.2, by_epoch=False, begin=0, end=7440),
+    dict(
+        type='CosineAnnealingLR', begin=7440, by_epoch=False, end=59520)
+]
+
+# runtime settings
+train_cfg = dict(type='IterBasedTrainLoop', max_iters=59520, val_interval=1488)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# default hooks - logger & checkpoint configs
+default_hooks = dict(
+
+    # print log every 100 iterations.
+    logger=dict(type='LoggerHook', interval=100, log_metric_by_epoch=False),
+
+    # enable the parameter scheduler.
+    param_scheduler=dict(type='ParamSchedulerHook'),
+
+    # save checkpoint every 5 epochs.
+    checkpoint=dict(type='CheckpointHook', by_epoch=False, interval=7440),
+)
+
+# tensorboard vis
+vis_backends = [dict(type='LocalVisBackend'),
+                dict(type='TensorboardVisBackend')]
+
+# data preprocessing
+norm_cfg = dict(type='SyncBN', requires_grad=True)
+crop_size = (512, 1024)
+data_preprocessor = dict(
+    type='SegDataPreProcessor',
+    mean=[0.0, 0.0, 0.0],
+    std=[1.0, 1.0, 1.0],
+    bgr_to_rgb=True,
+    pad_val=0,
+    seg_pad_val=255,
+    size=crop_size)
+
+model = dict(
+    type='EncoderDecoder',
+    backbone=dict(
+        type='hailoFPN',
+        depth=0.33,
+        width=0.125,
+        bb_channels_list=[128, 256, 512, 1024],
+        bb_num_repeats_list=[9, 15, 21, 12],
+        neck_channels_list=[256, 128, 128, 256, 256, 512],
+        neck_num_repeats_list=[9, 12, 12, 9]),
+    decode_head=dict(
+        type='PostProcess',
+        in_channels=16,
+        channels=128,
+        num_convs=1,
+        num_classes=10,
+        norm_cfg=norm_cfg,
+        align_corners=True,
+        loss_decode=dict(
+            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
+    # model training and testing settings
+    train_cfg=dict(),
+    test_cfg=dict(mode='whole'),
+    infer_wo_softmax=True)

configs/fcn/fcn_hailo_10classes.py

@@ -0,0 +1,75 @@
+# model settings
+_base_ = [
+    '../_base_/datasets/cityscapes10classes.py', '../_base_/default_runtime.py',
+]
+
+# optimizer
+optimizer = dict(type='Adam', lr=0.001, weight_decay=1e-5)
+optim_wrapper = dict(type='OptimWrapper', optimizer=optimizer, clip_grad=None)
+
+# learning policy
+param_scheduler = [
+	dict(
+		type='LinearLR', start_factor=0.2, by_epoch=False, begin=0, end=7440),
+    dict(
+        type='CosineAnnealingLR', begin=7440, by_epoch=False, end=59520)
+]
+
+# runtime settings
+train_cfg = dict(type='IterBasedTrainLoop', max_iters=59520, val_interval=1488)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# default hooks - logger & checkpoint configs
+default_hooks = dict(
+
+    # print log every 100 iterations.
+    logger=dict(type='LoggerHook', interval=100, log_metric_by_epoch=False),
+
+    # enable the parameter scheduler.
+    param_scheduler=dict(type='ParamSchedulerHook'),
+
+    # save checkpoint every 5 epochs.
+    checkpoint=dict(type='CheckpointHook', by_epoch=False, interval=7440),
+)
+
+# tensorboard vis
+vis_backends = [dict(type='LocalVisBackend'),
+                dict(type='TensorboardVisBackend')]
+
+# data preprocessing
+norm_cfg = dict(type='SyncBN', requires_grad=True)
+crop_size = (512, 1024)
+data_preprocessor = dict(
+    type='SegDataPreProcessor',
+    mean=[0.0, 0.0, 0.0],
+    std=[1.0, 1.0, 1.0],
+    bgr_to_rgb=True,
+    pad_val=0,
+    seg_pad_val=255,
+    size=crop_size)
+
+model = dict(
+    type='EncoderDecoder',
+    backbone=dict(
+        type='hailoFPN',
+        depth=0.33,
+        width=0.125,
+        bb_channels_list=[128, 256, 512, 1024],
+        bb_num_repeats_list=[9, 15, 21, 12],
+        neck_channels_list=[256, 128, 128, 256, 256, 512],
+        neck_num_repeats_list=[9, 12, 12, 9]),
+    decode_head=dict(
+        type='ConvHead',
+        in_channels=16,
+        channels=128,
+        num_convs=1,
+        num_classes=10,
+        norm_cfg=norm_cfg,
+        align_corners=True,
+        loss_decode=dict(
+            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
+    # model training and testing settings
+    train_cfg=dict(),
+    test_cfg=dict(mode='whole'),
+    infer_wo_softmax=True)

mmseg/datasets/__init__.py

@@ -4,6 +4,7 @@
 from .basesegdataset import BaseSegDataset
 from .chase_db1 import ChaseDB1Dataset
 from .cityscapes import CityscapesDataset
+from .cityscapes10classes import Cityscapes10classesDataset
 from .coco_stuff import COCOStuffDataset
 from .dark_zurich import DarkZurichDataset
 from .dataset_wrappers import MultiImageMixDataset
@@ -37,7 +38,7 @@
 # yapf: enable
 __all__ = [
     'BaseSegDataset', 'BioMedical3DRandomCrop', 'BioMedical3DRandomFlip',
-    'CityscapesDataset', 'PascalVOCDataset', 'ADE20KDataset',
+    'CityscapesDataset', 'Cityscapes10classesDataset', 'PascalVOCDataset', 'ADE20KDataset',
     'PascalContextDataset', 'PascalContextDataset59', 'ChaseDB1Dataset',
     'DRIVEDataset', 'HRFDataset', 'STAREDataset', 'DarkZurichDataset',
     'NightDrivingDataset', 'COCOStuffDataset', 'LoveDADataset',

mmseg/datasets/cityscapes10classes.py

@@ -0,0 +1,25 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmseg.registry import DATASETS
+from .basesegdataset import BaseSegDataset
+
+# [road, sidewalk, building, wall, vegetation, sky, car, person, traffic_light, truck]
+@DATASETS.register_module()
+class Cityscapes10classesDataset(BaseSegDataset):
+    """Cityscapes dataset.
+
+    The ``img_suffix`` is fixed to '_leftImg8bit.png' and ``seg_map_suffix`` is
+    fixed to '_gtFine_labelTrainIds.png' for Cityscapes dataset.
+    """
+    METAINFO = dict(
+        classes=('road', 'sidewalk', 'building', 'wall','vegetation',
+                 'sky', 'car', 'person', 'traffic light', 'truck'),
+        palette=[[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156],
+                 [190, 153, 153], [153, 153, 153], [250, 170,30], [220, 220, 0],
+                 [107, 142, 35], [152, 251, 152]])
+
+    def __init__(self,
+                 img_suffix='_leftImg8bit.png',
+                 seg_map_suffix='_gtFine_labelTrainIds.png',
+                 **kwargs) -> None:
+        super().__init__(
+            img_suffix=img_suffix, seg_map_suffix=seg_map_suffix, **kwargs)

tools/parse_cityscapes_to_10_classes.py

@@ -0,0 +1,32 @@
+import os
+from PIL import Image
+
+labels_list_dict = {0:0, 1:1, 2:2, 3:3, 8:4, 10:5, 13:6, 11:7, 6:8, 14:9}  # -> [road, sidewalk, building, wall, vegetation, sky, car, person, traffic_light, truck]
+
+def process_images(root_dir):
+    for root, _, files in os.walk(root_dir):
+        for filename in files:
+            if filename.endswith('labelTrainIds.png'):
+                image_path = os.path.join(root, filename)
+                try:
+                    image = Image.open(image_path)
+                    pixels = image.load()
+
+                    width, height = image.size
+                    for x in range(width):
+                        for y in range(height):
+                            pixel_value = pixels[x, y]
+                            if pixel_value not in labels_list_dict:
+                                pixels[x, y] = 255
+                            else:
+                                pixels[x,y] = labels_list_dict[pixels[x,y]]
+
+                    image.save(image_path)
+                    print(f"Processed: {image_path}")
+
+                except Exception as e:
+                    print(f"Error processing {image_path}: {e}")
+
+if __name__ == "__main__":
+    input_path = "/data/data/cityscapes10classes/gtFine/"
+    process_images(input_path)

tools/pytorch2onnx.py

@@ -19,15 +19,60 @@
 
 import torch.nn.functional as F
 
+from collections import OrderedDict
+import warnings
+
+def load_pretrained_weights_soft(model, checkpoint):
+
+    if 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    elif 'model' in checkpoint:
+        state_dict = checkpoint['model']
+    else:
+        state_dict = checkpoint
+
+    model_dict = model.state_dict()
+    new_state_dict = OrderedDict()
+    matched_layers, discarded_layers = [], []
+
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            k = k[7:] # discard module.
+
+        if k in model_dict and model_dict[k].size() == v.size():
+            new_state_dict[k] = v
+            matched_layers.append(k)
+        else:
+            discarded_layers.append(k)
+
+    model_dict.update(new_state_dict)
+    model.load_state_dict(model_dict)
+
+    if len(matched_layers) == 0:
+        warnings.warn(
+            'The pretrained weights "{}" cannot be loaded, '
+            'please check the key names manually '
+            '(** ignored and continue **)'
+        )
+    else:
+        print('Successfully loaded pretrained weights')
+        if len(discarded_layers) > 0:
+            print(
+                '** The following layers are discarded '
+                'due to unmatched keys or layer size: {}'.
+                format(discarded_layers)
+            )
+
 
 def parse_args():
     parser.add_argument('config', help='train config file path')
     parser.add_argument('--work-dir', help='the dir to save logs and models')
     parser.add_argument('--checkpoint', help='checkpoint file', default=None)
     parser.add_argument('--no_simplify', action='store_false')
-    parser.add_argument('--no_postprocess', action='store_false', default=False)
+    parser.add_argument('--no_postprocess', action='store_true', default=False)
     parser.add_argument('--shape', nargs=2, type=int, default=[1024, 1920])
     parser.add_argument('--out_name', default='fcn.onnx', type=str, help="Name for the onnx output")
+    parser.add_argument('--soft_weights_loading',action='store_true', default=False)
     parser.add_argument(
         '--cfg-options',
         nargs='+',
@@ -64,15 +109,13 @@ def __init__(self, model, args):
 
         def forward(self, x):
             x = self.model(x)
-            batch_size, C, H, W = x.shape
             if self.post_proc_flag:
-                    x = self.bilinear_resize(x)
-                    if C > 1:
-                        x = x.argmax(dim=1, keepdim=True)
+                x = self.bilinear_resize(x)
+                if x.shape[1] > 1:
+                    x = x.argmax(dim=1, keepdim=True)
             return x
 
 
-
 def main():
     args = parse_args()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -105,10 +148,13 @@ def main():
     model = runner.model
     if args.checkpoint:
         ckpt = torch.load(args.checkpoint, map_location='cpu')
-        if 'state_dict' in ckpt:
-            model.load_state_dict(ckpt['state_dict'])
+        if args.soft_weights_loading:
+            load_pretrained_weights_soft(model, ckpt)
         else:
-            model.load_state_dict(ckpt)
+            if 'state_dict' in ckpt:
+                model.load_state_dict(ckpt['state_dict'])
+            else:
+                model.load_state_dict(ckpt)
 
     # if repvgg style -> deploy
     for module in model.modules():