Merge branch 'master' into quickstart

Author: nickponline

libs/config.py

@@ -22,6 +22,10 @@
     (200, 130,   0) : 6,
 }
 
+LABELMAP_RGB = { k: (v[2], v[1], v[0]) for k, v in LABELMAP.items() }
+
+INV_LABELMAP_RGB = { v: k for k, v in LABELMAP_RGB.items() }
+
 train_ids = [
     "1d4fbe33f3_F1BE1D4184INSPIRE",
     "1df70e7340_4413A67E91INSPIRE",

libs/datasets.py

@@ -1,10 +1,3 @@
-import os
-from fastai.vision import *
-from fastai.callbacks.hooks import *
-from pathlib import PosixPath
-
-import numpy as np
-from libs.config import LABELS
 import libs.images2chips
 import sys
 import os
@@ -43,18 +36,3 @@ def download_dataset(dataset):
         libs.images2chips.run(dataset)
     else:
         print(f'chip folders "{image_chips}" and "{label_chips}" already exist, remove them to recreate chips.')
-
-def load_dataset(dataset, training_chip_size, bs):
-    """ Load a dataset, create batches and augmentation """
-
-    path = PosixPath(dataset)
-    label_path = path/'label-chips'
-    image_path = path/'image-chips'
-    image_files = get_image_files(image_path)
-    label_files = get_image_files(label_path)
-    get_y_fn = lambda x: label_path/f'{x.stem}{x.suffix}'
-    codes = np.array(LABELS)
-    src = SegmentationItemList.from_folder(image_path).split_by_fname_file('../valid.txt').label_from_func(get_y_fn, classes=codes)
-    # some data augmentation here
-    data = src.transform(get_transforms(flip_vert=True, max_warp=0., max_zoom=0., max_rotate=180.), size=training_chip_size, tfm_y=True).databunch(bs=bs)
-    return data

libs/datasets_fastai.py

@@ -0,0 +1,21 @@
+from fastai.vision import *
+from fastai.callbacks.hooks import *
+from pathlib import PosixPath
+
+import numpy as np
+from libs.config import LABELS
+
+def load_dataset(dataset, training_chip_size, bs):
+    """ Load a dataset, create batches and augmentation """
+
+    path = PosixPath(dataset)
+    label_path = path/'label-chips'
+    image_path = path/'image-chips'
+    image_files = get_image_files(image_path)
+    label_files = get_image_files(label_path)
+    get_y_fn = lambda x: label_path/f'{x.stem}{x.suffix}'
+    codes = np.array(LABELS)
+    src = SegmentationItemList.from_folder(image_path).split_by_fname_file('../valid.txt').label_from_func(get_y_fn, classes=codes)
+    # some data augmentation here
+    data = src.transform(get_transforms(flip_vert=True, max_warp=0., max_zoom=0., max_rotate=180.), size=training_chip_size, tfm_y=True).databunch(bs=bs)
+    return data

libs/datasets_keras.py

@@ -0,0 +1,65 @@
+from keras.preprocessing.image import ImageDataGenerator
+from keras.utils import Sequence, to_categorical
+from PIL import Image
+
+import numpy as np
+import random
+
+def load_dataset(dataset, bs, aug={'horizontal_flip': True, 'vertical_flip': True, 'rotation_range': 180}):
+    train_files = [f'{dataset}/image-chips/{fname}' for fname in load_lines(f'{dataset}/train.txt')]
+    valid_files = [f'{dataset}/image-chips/{fname}' for fname in load_lines(f'{dataset}/valid.txt')]
+    
+    train_seq = SegmentationSequence(
+        dataset,
+        train_files, 
+        ImageDataGenerator(**aug),
+        bs
+    )
+    
+    valid_seq = SegmentationSequence(
+        dataset,
+        valid_files,
+        ImageDataGenerator(), # don't augment validation set
+        bs
+    )
+    
+    return train_seq, valid_seq
+
+def load_lines(fname):
+    with open(fname, 'r') as f:
+        return [l.strip() for l in f.readlines()]
+
+def load_img(fname):
+    return np.array(Image.open(fname))
+
+def mask_to_classes(mask):
+    return to_categorical(mask[:,:,0], 6)
+
+class SegmentationSequence(Sequence):
+    def __init__(self, dataset, image_files, datagen, bs):
+        self.label_path = f'{dataset}/label-chips'
+        self.image_path = f'{dataset}/image-chips'
+        self.image_files = image_files
+        random.shuffle(self.image_files)
+
+        self.datagen = datagen
+        self.bs = bs
+
+    def __len__(self):
+        return int(np.ceil(len(self.image_files) / float(self.bs)))
+
+    def __getitem__(self, idx):
+        image_files = self.image_files[idx*self.bs:(idx+1)*self.bs]
+        label_files = [fname.replace(self.image_path, self.label_path) for fname in image_files]
+
+        images = [load_img(fname) for fname in image_files]
+        labels = [mask_to_classes(load_img(fname)) for fname in label_files]
+
+        ts = [self.datagen.get_random_transform(im.shape) for im in images]
+        images = [self.datagen.apply_transform(im, ts) for im, ts in zip(images, ts)]
+        labels = [self.datagen.apply_transform(im, ts) for im, ts in zip(labels, ts)]
+
+        return np.array(images), np.array(labels)
+
+    def on_epoch_end(self):
+        random.shuffle(self.image_files)

libs/inference_keras.py

@@ -0,0 +1,71 @@
+from PIL import Image
+import numpy as np
+import math
+from keras import models
+import os
+
+from libs.config import train_ids, test_ids, val_ids, LABELMAP_RGB
+
+def category2mask(img):
+    """ Convert a category image to color mask """
+    if len(img) == 3:
+        if img.shape[2] == 3:
+            img = img[:, :, 0]
+
+    mask = np.zeros(img.shape[:2] + (3, ), dtype='uint8')
+
+    for category, mask_color in LABELMAP_RGB.items():
+        locs = np.where(img == category)
+        mask[locs] = mask_color
+
+    return mask
+
+def chips_from_image(img, size=300):
+    shape = img.shape
+
+    chip_count = math.ceil(shape[1] / size) * math.ceil(shape[0] / size)
+
+    chips = []
+    for x in range(0, shape[1], size):
+        for y in range(0, shape[0], size):
+            chip = img[y:y+size, x:x+size, :]
+            y_pad = size - chip.shape[0]
+            x_pad = size - chip.shape[1]
+            chip = np.pad(chip, [(0, y_pad), (0, x_pad), (0, 0)], mode='constant')
+            chips.append((chip, x, y))
+    return chips
+
+def run_inference_on_file(imagefile, predsfile, model, size=300):
+    with Image.open(imagefile).convert('RGB') as img:
+        nimg = np.array(Image.open(imagefile).convert('RGB'))
+        shape = nimg.shape
+        chips = chips_from_image(nimg)
+
+    chips = [(chip, xi, yi) for chip, xi, yi in chips if chip.sum() > 0]
+    prediction = np.zeros(shape[:2], dtype='uint8')
+    chip_preds = model.predict(np.array([chip for chip, _, _ in chips]), verbose=True)
+
+    for (chip, x, y), pred in zip(chips, chip_preds):
+        category_chip = np.argmax(pred, axis=-1) + 1
+        section = prediction[y:y+size, x:x+size].shape
+        prediction[y:y+size, x:x+size] = category_chip[:section[0], :section[1]]
+
+    mask = category2mask(prediction)
+    Image.fromarray(mask).save(predsfile)
+
+def run_inference(dataset, model=None, model_path=None, basedir='.'):
+    if model is None and model_path is None:
+        raise Exception("model or model_path required")
+
+    if model is None:
+        model = models.load_model(model_path)
+
+    for scene in train_ids + val_ids + test_ids:
+        imagefile = f'{dataset}/images/{scene}-ortho.tif'
+        predsfile = os.path.join(basedir, f'{scene}-prediction.png')
+
+        if not os.path.exists(imagefile):
+            continue
+
+        print(f'running inference on image {imagefile}.')
+        run_inference_on_file(imagefile, predsfile, model)

libs/models_keras.py

@@ -0,0 +1,128 @@
+from keras import layers, models
+import numpy as np
+import tensorflow as tf
+
+def build_unet(size=300, basef=64, maxf=512, encoder='resnet50', pretrained=True):
+    input = layers.Input((size, size, 3))
+
+    encoder_model = make_encoder(input, name=encoder, pretrained=pretrained)
+
+    crosses = []
+
+    for layer in encoder_model.layers:
+        # don't end on padding layers
+        if type(layer) == layers.ZeroPadding2D:
+            continue
+        idx = get_scale_index(size, layer.output_shape[1])
+        if idx is None:
+            continue
+        if idx >= len(crosses):
+            crosses.append(layer)
+        else:
+            crosses[idx] = layer
+
+    x = crosses[-1].output
+    for scale in range(len(crosses)-2, -1, -1):
+        nf = min(basef * 2**scale, maxf)
+        x = upscale(x, nf)
+        x = act(x)
+        x = layers.Concatenate()([
+            pad_to_scale(x, scale, size=size),
+            pad_to_scale(crosses[scale].output, scale, size=size)
+        ])
+        x = conv(x, nf)
+        x = act(x)
+
+    x = conv(x, 6)
+    x = layers.Activation('softmax')(x)
+
+    return models.Model(input, x)
+
+def make_encoder(input, name='resnet50', pretrained=True):
+    if name == 'resnet18':
+        from classification_models.keras import Classifiers
+        ResNet18, _ = Classifiers.get('resnet18')
+        model = ResNet18(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    elif name == 'resnet50':
+        from keras.applications.resnet import ResNet50
+        model = ResNet50(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    elif name == 'resnet101':
+        from keras.applications.resnet import ResNet101
+        model = ResNet101(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    elif name == 'resnet152':
+        from keras.applications.resnet import ResNet152
+        model = ResNet152(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    elif name == 'vgg16':
+        from keras.applications.vgg16 import VGG16
+        model = VGG16(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    elif name == 'vgg19':
+        from keras.applications.vgg19 import VGG19
+        model = VGG19(
+            weights='imagenet' if pretrained else None,
+            input_tensor=input,
+            include_top=False
+        )
+    else:
+        raise Exception(f'unknown encoder {name}')
+
+    return model
+
+def get_scale_index(in_size, l_size):
+    for i in range(8):
+        s_size = in_size // (2 ** i)
+        if abs(l_size - s_size) <= 4:
+            return i
+    return None
+
+def pad_to_scale(x, scale, size=300):
+    expected = int(np.ceil(size / (2. ** scale)))
+    diff = expected - int(x.shape[1])
+    if diff > 0:
+        left = diff // 2
+        right = diff - left
+        x = reflectpad(x, (left, right))
+    elif diff < 0:
+        left = -diff // 2
+        right = -diff - left
+        x = layers.Cropping2D(((left, right), (left, right)))(x)
+    return x
+
+def reflectpad(x, pad):
+    return layers.Lambda(lambda x: tf.pad(x, [(0, 0), pad, pad, (0, 0)], 'REFLECT'))(x)
+
+def upscale(x, nf):
+    x = layers.UpSampling2D((2, 2))(x)
+    x = conv(x, nf, kernel_size=(1, 1))
+    return x
+
+def act(x):
+    x = layers.BatchNormalization()(x)
+    x = layers.LeakyReLU(0.2)(x)
+    return x
+
+def conv(x, nf, kernel_size=(3, 3), **kwargs):
+    padleft = (kernel_size[0] - 1) // 2
+    padright = kernel_size[0] - 1 - padleft
+    if padleft > 0 or padright > 0:
+        x = reflectpad(x, (padleft, padright))
+    return layers.Conv2D(nf, kernel_size=kernel_size, padding='valid', **kwargs)(x)

libs/scoring.py

@@ -52,10 +52,11 @@ def plot_confusion_matrix(y_true, y_pred, classes,
     im = ax.imshow(cm, interpolation='nearest', cmap=cmap)
     ax.figure.colorbar(im, ax=ax)
 
+    base, fname = os.path.split(title)
     ax.set(xticks=np.arange(cm.shape[1]),
            yticks=np.arange(cm.shape[0]),
            xticklabels=classes, yticklabels=classes,
-           title=title,
+           title=fname,
            ylabel='True label',
            xlabel='Predicted label')
 
@@ -79,6 +80,7 @@ def plot_confusion_matrix(y_true, y_pred, classes,
     if not os.path.isdir(savedir):
         os.mkdir(savedir)
     savefile = savedir + '/score-' + title
+
     plt.savefig(savefile)
     return savefile, cm
 
@@ -117,7 +119,7 @@ def score_masks(labelfile, predictionfile):
 
     return precision, recall, f1, savefile
 
-def score_predictions(dataset):
+def score_predictions(dataset, basedir='.'):
 
     scores = []
 
@@ -132,7 +134,7 @@ def score_predictions(dataset):
     for scene in test_ids:
 
         labelfile = f'{dataset}/labels/{scene}-label.png'
-        predsfile = f"{scene}-prediction.png"
+        predsfile = os.path.join(basedir, f"{scene}-prediction.png")
 
         if not os.path.exists(labelfile):
             continue

libs/training.py

@@ -8,7 +8,7 @@
 from libs import inference
 from libs import scoring
 from libs.util import MySaveModelCallback, ExportCallback, MyCSVLogger, Precision, Recall, FBeta
-from libs import datasets
+from libs import datasets_fastai
 
 import wandb
 from wandb.fastai import WandbCallback
@@ -41,7 +41,7 @@ def train_model(dataset):
         FBeta(average='weighted', beta=1, clas_idx=1),
     ]
 
-    data = datasets.load_dataset(dataset, size, bs)
+    data = datasets_fastai.load_dataset(dataset, size, bs)
     encoder_model = models.resnet18
     learn = unet_learner(data, encoder_model, path='models', metrics=metrics, wd=wd, bottle=True, pretrained=pretrained)