lmnet: Add cityscapes loader (#300)

Author: Joeper214

lmnet/configs/core/segmentation/segnet_quantize_cityscapes.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Blueoil Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+from easydict import EasyDict
+import tensorflow as tf
+
+from lmnet.common import Tasks
+from lmnet.networks.segmentation.lm_segnet_quantize import LmSegnetQuantize
+from lmnet.datasets.cityscapes import Cityscapes
+from lmnet.data_processor import Sequence
+from lmnet.pre_processor import (
+    DivideBy255,
+    Resize
+)
+from lmnet.data_augmentor import (
+    Brightness,
+    Color,
+    Contrast,
+    FlipLeftRight,
+    Hue,
+)
+from lmnet.quantizations import (
+    binary_mean_scaling_quantizer,
+    linear_mid_tread_half_quantizer,
+)
+
+IS_DEBUG = False
+
+NETWORK_CLASS = LmSegnetQuantize
+DATASET_CLASS = Cityscapes
+
+IMAGE_SIZE = [160, 320]
+BATCH_SIZE = 8
+DATA_FORMAT = "NHWC"
+TASK = Tasks.SEMANTIC_SEGMENTATION
+CLASSES = DATASET_CLASS.classes
+
+MAX_STEPS = 150000
+SAVE_STEPS = 3000
+TEST_STEPS = 1000
+SUMMARISE_STEPS = 1000
+
+# distributed training
+IS_DISTRIBUTION = False
+
+# pretrain
+IS_PRETRAIN = False
+PRETRAIN_VARS = []
+PRETRAIN_DIR = ""
+PRETRAIN_FILE = ""
+
+# for debug
+# BATCH_SIZE = 2
+# SUMMARISE_STEPS = 1
+# IS_DEBUG = True
+
+PRE_PROCESSOR = Sequence([
+    Resize(size=IMAGE_SIZE),
+    DivideBy255()
+])
+POST_PROCESSOR = None
+
+NETWORK = EasyDict()
+NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
+NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
+NETWORK.IMAGE_SIZE = IMAGE_SIZE
+NETWORK.BATCH_SIZE = BATCH_SIZE
+NETWORK.DATA_FORMAT = DATA_FORMAT
+NETWORK.ACTIVATION_QUANTIZER = linear_mid_tread_half_quantizer
+NETWORK.ACTIVATION_QUANTIZER_KWARGS = {
+    'bit': 2,
+    'max_value': 2
+}
+NETWORK.WEIGHT_QUANTIZER = binary_mean_scaling_quantizer
+NETWORK.WEIGHT_QUANTIZER_KWARGS = {}
+
+DATASET = EasyDict()
+DATASET.BATCH_SIZE = BATCH_SIZE
+DATASET.DATA_FORMAT = DATA_FORMAT
+DATASET.PRE_PROCESSOR = PRE_PROCESSOR
+DATASET.AUGMENTOR = Sequence([
+    Brightness((0.75, 1.25)),
+    Color((0.75, 1.25)),
+    Contrast((0.75, 1.25)),
+    FlipLeftRight(),
+    Hue((-10, 10)),
+])
+DATASET.ENABLE_PREFETCH = True

lmnet/lmnet/datasets/cityscapes.py

@@ -0,0 +1,148 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Blueoil Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+# Support for cityscapes dataset
+# https://www.cityscapes-dataset.com/
+
+import functools
+import glob
+import os.path
+import numpy as np
+
+from PIL import Image
+
+from lmnet.datasets.base import SegmentationBase
+
+
+class Cityscapes(SegmentationBase):
+    available_subsets = ["train", "validation", "test"]
+    extend_dir = "cityscapes"
+    classes = [
+        "unlabeled",
+        "ego vehicle",
+        "rectification boarder",
+        "out of roi",
+        "static",
+        "dynamic",
+        "ground",
+        "road",
+        "sidewalk",
+        "parking",
+        "rail track",
+        "building",
+        "wall",
+        "fence",
+        "guard rail",
+        "bridge",
+        "tunnel",
+        "pole",
+        "polegroup",
+        "traffic light",
+        "traffic sign",
+        "vegetation",
+        "terrain",
+        "sky",
+        "person",
+        "rider",
+        "car",
+        "truck",
+        "bus",
+        "caravan",
+        "trailer",
+        "train",
+        "motorcycle",
+        "bicycle",
+    ]
+    num_classes = len(classes)
+
+    def __init__(self, batch_size=10, *args, **kwargs):
+        super().__init__(batch_size=batch_size, *args, **kwargs)
+
+    @property
+    def label_colors(self):
+        unlabeled = [0, 0, 0]
+        ego_vehicle = [0, 0, 0]
+        rectification_boarder = [0, 0, 0]
+        out_of_roi = [0, 0, 0]
+        static = [0, 0, 0]
+        dynamic = [111, 74, 0]
+        ground = [81, 0, 81]
+        road = [128, 64, 128]
+        sidewalk = [244, 35, 232]
+        parking = [250, 170, 160]
+        rail_track = [230, 150, 140]
+        building = [70, 70, 70]
+        wall = [102, 102, 156]
+        fence = [190, 153, 153]
+        guard_rail = [180, 165, 180]
+        bridge = [150, 100, 100]
+        tunnel = [150, 120, 90]
+        pole = [153, 153, 153]
+        polegroup = [153, 153, 153]
+        traffic_light = [250, 170, 30]
+        traffic_sign = [220, 220, 0]
+        vegetation = [107, 142, 35]
+        terrain = [152, 251, 152]
+        sky = [70, 130, 180]
+        person = [220, 20, 60]
+        rider = [255, 0, 0]
+        car = [0, 0, 142]
+        truck = [0, 0, 70]
+        bus = [0, 60, 100]
+        caravan = [0, 0, 90]
+        trailer = [0, 0, 110]
+        train = [0, 80, 100]
+        motorcycle = [0, 0, 230]
+        bicycle = [119, 11, 32]
+
+        return np.array([
+            unlabeled, ego_vehicle, rectification_boarder, out_of_roi, static,
+            dynamic, ground, road, sidewalk, parking, rail_track, building,
+            wall, fence, guard_rail, bridge, tunnel, pole, polegroup,
+            traffic_light, traffic_sign, vegetation, terrain, sky, person,
+            rider, car, truck, bus, caravan, trailer, train, motorcycle,
+            bicycle])
+
+    @functools.lru_cache(maxsize=None)
+    def files_and_annotations(self):
+        split = "train"
+        if self.subset == "validation":
+            split = "val"
+        elif self.subset == "test":
+            split = "test"
+        polygons_json = glob.glob(os.path.join(self.data_dir, "gtFine", split, "*", "*_gt*_polygons.json"))
+        polygons_json.sort()
+
+        labelIds = [i.replace("_polygons.json", "_labelIds.png") for i in polygons_json]
+        leftImg8bit = [i.replace(
+            os.path.join(self.data_dir, "gtFine"),
+            os.path.join(self.data_dir, "leftImg8bit")
+        ).replace("_gtFine_polygons.json", "_leftImg8bit.png") for i in polygons_json]
+
+        return leftImg8bit, labelIds
+
+    def __getitem__(self, i):
+        imgs, labels = self.files_and_annotations()
+        img = Image.open(imgs[i])
+        label = Image.open(labels[i])
+
+        return np.array(img), np.array(label)
+
+    def __len__(self):
+        return len(self.files_and_annotations()[0])
+
+    @property
+    def num_per_epoch(self):
+        return len(self.files_and_annotations()[0])

lmnet/tests/check_dataset_storage.py

@@ -37,6 +37,7 @@
 from lmnet.datasets.caltech101 import Caltech101
 from lmnet.datasets.cifar10 import Cifar10
 from lmnet.datasets.camvid import Camvid
+from lmnet.datasets.cityscapes import Cityscapes
 from lmnet.datasets.ilsvrc_2012 import Ilsvrc2012
 from lmnet.datasets.lm_things_on_a_table import LmThingsOnATable
 from lmnet.datasets.mscoco import (
@@ -558,6 +559,38 @@ def test_camvid():
         assert labels.shape[2] == image_size[1]
 
 
+def test_cityscapes():
+    batch_size = 3
+    image_size = [256, 512]
+    dataset = Cityscapes(
+        batch_size=batch_size,
+        pre_processor=Resize(image_size))
+    dataset = DatasetIterator(dataset)
+
+    val_dataset = Cityscapes(
+        subset="validation",
+        batch_size=batch_size,
+        pre_precessor=Resize(image_size)
+    )
+
+    assert dataset.num_classes == 34
+    assert dataset.num_per_epoch == 2975
+    assert val_dataset.num_per_epoch == 500
+
+    for _ in range(STEP_SIZE):
+        images, labels = dataset.feed()
+        assert isinstance(images, np.ndarray)
+        assert images.shape[0] == batch_size
+        assert images.shape[1] == image_size[0]
+        assert images.shape[2] == image_size[1]
+        assert images.shape[3] == 3
+
+        assert isinstance(labels, np.ndarray)
+        assert labels.shape[0] == batch_size
+        assert labels.shape[1] == image_size[0]
+        assert labels.shape[2] == image_size[1]
+
+
 def test_lm_things_of_a_table():
     batch_size = 3
     image_size = [256, 512]
@@ -862,6 +895,7 @@ def test_bdd100k():
     test_caltech101()
     test_cifar10()
     test_camvid()
+    test_cityscapes()
     test_pascalvoc_2007()
     test_pascalvoc_2007_not_skip_difficult()
     test_pascalvoc_2007_with_target_classes()