Merge pull request #1 from dansola/light_weight

Light weight

Author: dansola

README.md

@@ -26,7 +26,7 @@ These proposals are incorporated into a Dataset class where a random index dicti
 Then PGA can be used on an image and its corresponding proposal dictionaries.
 
 ```python
-from models.basic_pga.basic_pga_parts import BlockPGA
+from src.models.basic_pga.basic_pga_parts import BlockPGA
 import torch
 
 img = torch.rand(1,3,500,500) # test image

notebooks/__init__.py

@@ -4,13 +4,13 @@ kiwisolver==1.3.1
 matplotlib==3.3.3
 olefile==0.46
 opencv-python==4.5.1.48
-Pillow @ file:///tmp/build/80754af9/pillow_1609786786540/work
+Pillow==8.1.0
 pyparsing==2.4.7
 python-dateutil==2.8.1
-six @ file:///tmp/build/80754af9/six_1605205327372/work
+six==1.15.0
 torch==1.7.1
 torchvision==0.8.2
-typing-extensions @ file:///tmp/build/80754af9/typing_extensions_1598376058250/work
+typing-extensions
 numpy~=1.19.2
 pillow~=8.1.0
 wandb~=0.10.15

requirements.txt

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+import os
+import sys
+
+from setuptools import find_packages, setup
+from setuptools.command.install import install
+
+from src import VERSION
+
+
+class VerifyVersionCommand(install):
+    """Custom command to verify that the git tag matches our version"""
+
+    description = "verify that the package git tag matches our version"
+
+    def run(self):
+        tag = os.getenv("COMMIT_TAG")
+
+        if tag != VERSION:
+            info = "Git tag: {0} does not match the version of this app: {1}".format(
+                tag, VERSION
+            )
+            sys.exit(info)
+
+
+with open("requirements.txt") as f:
+    DEPENDENCIES = f.read().splitlines()
+
+setup(
+    name="src",
+    packages=find_packages(),
+    version=VERSION,
+    description="Light weight semantic segmentation.",
+    author="Daniel Sola",
+    license="MIT",
+    install_requires=DEPENDENCIES,
+    python_requires=">=3.8",
+    url="https://github.com/dansola/PGA-Net",
+    cmdclass={"verify": VerifyVersionCommand},
+)

setup.py

@@ -0,0 +1 @@
+VERSION = "0.0.1"

src/__init__.py

@@ -0,0 +1,268 @@
+import os
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+import skimage.transform
+
+from torch.utils import data
+
+from datasets.utils import recursive_glob, Compose, RandomHorizontallyFlip, RandomRotate, Scale
+
+
+class City(data.Dataset):
+    """cityscapesLoader
+    https://www.cityscapes-dataset.com
+    Data is derived from CityScapes, and can be downloaded from here:
+    https://www.cityscapes-dataset.com/downloads/
+    Many Thanks to @fvisin for the loader repo:
+    https://github.com/fvisin/dataset_loaders/blob/master/dataset_loaders/images/cityscapes.py
+    """
+
+    colors = [  # [  0,   0,   0],
+        [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],
+        [0, 130, 180],
+        [220, 20, 60],
+        [255, 0, 0],
+        [0, 0, 142],
+        [0, 0, 70],
+        [0, 60, 100],
+        [0, 80, 100],
+        [0, 0, 230],
+        [119, 11, 32],
+    ]
+
+    label_colours = dict(zip(range(19), colors))
+
+    mean_rgb = {
+        "pascal": [103.939, 116.779, 123.68],
+        "cityscapes": [0.0, 0.0, 0.0],
+    }  # pascal mean for PSPNet and ICNet pre-trained model
+
+    def __init__(
+            self,
+            root,
+            split="train",
+            is_transform=False,
+            img_size=(256, 512),
+            augmentations=None,
+            img_norm=True,
+            version="cityscapes",
+            test_mode=False,
+    ):
+        """__init__
+        :param root:
+        :param split:
+        :param is_transform:
+        :param img_size:
+        :param augmentations
+        """
+        self.root = root
+        self.split = split
+        self.is_transform = is_transform
+        self.augmentations = augmentations
+        self.img_norm = img_norm
+        self.n_classes = 19
+        self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
+        self.mean = np.array(self.mean_rgb[version])
+        self.files = {}
+
+        self.images_base = os.path.join(self.root, "leftImg8bit", self.split)
+        self.annotations_base = os.path.join(self.root, "gtFine", self.split)
+
+        self.files[split] = recursive_glob(rootdir=self.images_base, suffix=".png")
+
+        self.void_classes = [0, 1, 2, 3, 4, 5, 6, 9, 10, 14, 15, 16, 18, 29, 30, -1]
+        self.valid_classes = [
+            7,
+            8,
+            11,
+            12,
+            13,
+            17,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            31,
+            32,
+            33,
+        ]
+        self.class_names = [
+            "unlabelled",
+            "road",
+            "sidewalk",
+            "building",
+            "wall",
+            "fence",
+            "pole",
+            "traffic_light",
+            "traffic_sign",
+            "vegetation",
+            "terrain",
+            "sky",
+            "person",
+            "rider",
+            "car",
+            "truck",
+            "bus",
+            "train",
+            "motorcycle",
+            "bicycle",
+        ]
+
+        self.ignore_index = 250
+        self.class_map = dict(zip(self.valid_classes, range(19)))
+
+        if not self.files[split]:
+            raise Exception("No files for split=[%s] found in %s" % (split, self.images_base))
+
+        print("Found %d %s images" % (len(self.files[split]), split))
+
+    def __len__(self):
+        """__len__"""
+        return len(self.files[self.split])
+
+    def __getitem__(self, index):
+        """__getitem__
+        :param index:
+        """
+        img_path = self.files[self.split][index].rstrip()
+        lbl_path = os.path.join(
+            self.annotations_base,
+            img_path.split(os.sep)[-2],
+            os.path.basename(img_path)[:-15] + "gtFine_labelIds.png",
+        )
+
+        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = np.float32(img)
+
+        lbl = cv2.imread(lbl_path, cv2.IMREAD_GRAYSCALE)  # GRAY 1 channel ndarray with shape H * W
+        lbl = self.encode_segmap(np.array(lbl, dtype=np.uint8))
+
+        # img = np.array(Image.open(img_path))
+        # img = np.array(img, dtype=np.uint8)
+        #
+        # lbl = np.array(Image.open(lbl_path))
+        # lbl = self.encode_segmap(np.array(lbl, dtype=np.uint8))
+
+        if self.augmentations is not None:
+            img, lbl = self.augmentations(img, lbl)
+
+        if self.is_transform:
+            img, lbl = self.transform(img, lbl)
+
+        return {
+            'image': img,
+            'mask': lbl
+        }
+
+    def transform(self, img, lbl):
+        """transform
+        :param img:
+        :param lbl:
+        """
+        img = skimage.transform.resize(img, (self.img_size[0], self.img_size[1]),
+                                       mode='edge',
+                                       anti_aliasing=False,
+                                       anti_aliasing_sigma=None,
+                                       preserve_range=True,
+                                       order=0)  # uint8 with RGB mode
+        img = img[:, :, ::-1]  # RGB -> BGR
+        img = img.astype(np.float64)
+        img -= self.mean
+        if self.img_norm:
+            # Resize scales images from 0 to 255, thus we need
+            # to divide by 255.0
+            img = img.astype(float) / 255.0
+        # NHWC -> NCHW
+        img = img.transpose(2, 0, 1)
+
+        classes = np.unique(lbl)
+        lbl = lbl.astype(float)
+        lbl = skimage.transform.resize(lbl, (self.img_size[0], self.img_size[1]),
+                                       mode='edge',
+                                       anti_aliasing=False,
+                                       anti_aliasing_sigma=None,
+                                       preserve_range=True,
+                                       order=0)
+        lbl = lbl.astype(int)
+
+        if not np.all(classes == np.unique(lbl)):
+            print("WARN: resizing labels yielded fewer classes")
+
+        if not np.all(np.unique(lbl[lbl != self.ignore_index]) < self.n_classes):
+            print("after det", classes, np.unique(lbl))
+            raise ValueError("Segmentation map contained invalid class values")
+
+        img = torch.from_numpy(img).float()
+        lbl = torch.from_numpy(lbl).long()
+
+        return img, lbl
+
+    def decode_segmap(self, temp):
+        r = temp.copy()
+        g = temp.copy()
+        b = temp.copy()
+        for l in range(0, self.n_classes):
+            r[temp == l] = self.label_colours[l][0]
+            g[temp == l] = self.label_colours[l][1]
+            b[temp == l] = self.label_colours[l][2]
+
+        rgb = np.zeros((temp.shape[0], temp.shape[1], 3))
+        rgb[:, :, 0] = r / 255.0
+        rgb[:, :, 1] = g / 255.0
+        rgb[:, :, 2] = b / 255.0
+        return rgb
+
+    def encode_segmap(self, mask):
+        # Put all void classes to zero
+        for _voidc in self.void_classes:
+            mask[mask == _voidc] = self.ignore_index
+        for _validc in self.valid_classes:
+            mask[mask == _validc] = self.class_map[_validc]
+        return mask
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+    augmentations = Compose([Scale(2048), RandomRotate(10), RandomHorizontallyFlip(0.5)])
+
+    local_path = "/datasets01/cityscapes/112817/"
+    dst = cityscapesLoader(local_path, is_transform=True, augmentations=augmentations)
+    bs = 4
+    trainloader = data.DataLoader(dst, batch_size=bs, num_workers=0)
+    for i, data_samples in enumerate(trainloader):
+        imgs, labels = data_samples
+        import pdb
+
+        pdb.set_trace()
+        imgs = imgs.numpy()[:, ::-1, :, :]
+        imgs = np.transpose(imgs, [0, 2, 3, 1])
+        f, axarr = plt.subplots(bs, 2)
+        for j in range(bs):
+            axarr[j][0].imshow(imgs[j])
+            axarr[j][1].imshow(dst.decode_segmap(labels.numpy()[j]))
+        plt.show()
+        a = input()
+        if a == "ex":
+            break
+        else:
+            plt.close()