Writing a custom transformer and unit testing it


#1

I’m trying to write a custom transformer, and unit test it. I was wondering if someone could help me get this working.

This is the transformer code for performing image alignment:

import cv2
import numpy as np

from sklearn.externals import joblib
from step.base import BaseTransformer

import common.logger as logger

logger = logger.init_logger('step', 'experiment/step.log')


class ImageAlignmentTransformer(BaseTransformer):
    def __init__(self):
        self.aligner = align

    def fit(self, img_ref, img_unalign):
        return self

    def transform(self, img_ref, img_unalign):
        img_align = self.aligner(img_ref, img_unalign)
        return {'img_align': img_align}

    def save(self, filepath):
        joblib.dump(self.aligner, filepath)

    def load(self, filepath):
        self.aligner = joblib.load(filepath)
        return self


def align(img_ref, img_unalign):
    """
    Align an image using a reference image.
    :param img_ref: Reference image used to compute warp matrix for aligning images.
    :param img_unalign: Unaligned input image.
    :return: img_align: Aligned image.
    """
    try:
        p1 = img_ref[300:1900, 300:2200, 0].astype(np.float32)
        p2 = img_unalign[300:1900, 300:2200, 0].astype(np.float32)
    except:
        logger.warn("Can't extract patch, falling back to whole image")
        p1 = img_ref[:, :, 0]
        p2 = img_unalign[:, :, 0]

    warp_mode = cv2.MOTION_EUCLIDEAN
    warp_matrix = np.eye(2, 3, dtype=np.float32)
    criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 1000, 1e-7)
    (cc, warp_matrix) = cv2.findTransformECC(p1, p2, warp_matrix, warp_mode, criteria)
    logger.debug("align: cc:{}".format(cc))

    # TODO: Remove the border reflect option, after creating a separate step for padding images.
    img_align = cv2.warpAffine(img_unalign, warp_matrix, (img_ref.shape[1], img_ref.shape[0]),
                               flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP,
                               borderMode=cv2.BORDER_REFLECT_101,
                               borderValue=0)
    img_align[img_align == 0] = np.average(img_align)

    return img_align

and this is the unit test

from unittest import TestCase

import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
import tifffile as tiff

import common.logger as logger
from step.base import Step
from step.preprocessing.image_alignment import align, ImageAlignmentTransformer

class TestImageAlignment(TestCase):
    def setUp(self):
        self.logger = logger.init_logger('test', 'output/test-image-alignment.log')

        # current working director
        cwd = os.getcwd()

        # dataset location
        self.dataset_dir = '{}/../../dataset/kaggle/dstl-sifd'.format(cwd)
        #self.dataset_dir = './dataset'

        # dataset metadata
        self.dataset_metadata_filename = '{}/../../dataset/kaggle/dstl-sifd/train_wkt_v4.csv'.format(cwd)
        #self.dataset_metadata_filename = './dataset/train_wkt_v4.csv'

        # create output directories
        self.output_dir = './output/image'
        create_dir(self.output_dir)

    def test_image_alignment_step(self):
        self.logger.info("Loading images")

        # read a sample
        sample_id = '6120_2_2'  # 6140_3_1, 6120_2_2, 6070_2_2

        # read rgb image
        img_rgb = tiff.imread('{}/three_band/{}.tif'.format(self.dataset_dir, sample_id)).transpose([1, 2, 0])
        nw, nh, _ = img_rgb.shape
        self.logger.info('rgb shape: {}'.format(img_rgb.shape))

        # read a-band image, low resolution
        img_a = tiff.imread("{}/sixteen_band/{}_A.tif".format(self.dataset_dir, sample_id)).transpose([1, 2, 0])
        self.logger.info('a-band image shape: {}'.format(img_a.shape))

        # upscale a-band image using opencv, note cv2.resize dsize is (nh, nw)
        rescaled_a = cv2.resize(img_a, (nh, nw), interpolation=cv2.INTER_CUBIC)
        self.logger.info('rescaled a-band image shape: {}'.format(rescaled_a.shape))

        data_train = {'input':
                        {
                            'img_ref': img_rgb,
                            'img_unalign': rescaled_a,
                        }
                     }

        align_step = Step(name='Image Aligner',
                          transformer=ImageAlignmentTransformer(),
                          input_data=[data_train['input']],
                          # adapter={'img_ref'     : ([('input', 'img_ref')]),
                          #          'img_unalign' : ([('input', 'img_unalign')])},
                          cache_dirpath='./cache')

        output = align_step.transform(data_train)

        print(output.shape)


def create_dir(path):
        dir = os.path.dirname(path)
        if not os.path.exists(dir):
                os.makedirs(dir)

I get the following error:

  File "/project/geospatial/application/cs230-sifd/test/step/preprocessing/test_image_alignment.py", line 107, in test_image_alignment_step
    output = align_step.transform(data_train)
  File "/project/geospatial/application/cs230-sifd/source/step/base.py", line 262, in transform
    step_inputs[input_data_part] = data[input_data_part]
TypeError: unhashable type: 'dict'

How should I feed the data input to this single step?


#2

Hi,

It seems you’re trying to pass a dict object as a key to another dict - the error can be reproduces by following code:

first = {"a":"b"}
second = {}
second[first] = "value"

I’m not familiar with code in step package, but my guess is that you could try one of those:

  • change input_data in Step constructor to [data_train] or data_train and uncomment adapter parameter
  • change input_data in Step constructor to single image data, like
    [data_train['input']['img_ref'], data_train['input']['img_unalign']]
  • change data_train in .transform() to match what you put in Step constructor, like data_train["input"] or [data_train["input"]]

If this doesn’t help, I’ll contact someone more familiar with this code and we’ll try to find a better solution tomorrow morning


#3

Hi,
You should actually pass just the key ‘input’ i.e.:

input_data=['input'],

The idea behind it is that you pass just one data dictionary in your case data_train but you could have different things in it. For instance you need ‘input’ for tabular data ‘input_text’ for textual and ‘input_image’ for image data in your pipeline.
However we do realize that it could be (and is) confusing and are working on improving the API as we speak.

By the way we should have steps released as a package (with improvements and documentation) in the next 2 weeks!


#4

Hello Jakub,

Thanks for the reply.

I’m still not clear on what my method signatures should look like. Would you mind editing the following code snippet to show me how the method signatures should look like for the constructor, fit and transform?

class ImageAlignmentTransformer(BaseTransformer):
    def __init__(self):
        self.aligner = align

    def fit(self, img_ref, img_unalign):
        return self

    def transform(self, img_ref, img_unalign):
        img_align = self.aligner(img_ref, img_unalign)
        return {'img_align': img_align}

    def save(self, filepath):
        joblib.dump(self.aligner, filepath)

    def load(self, filepath):
        self.aligner = joblib.load(filepath)
        return self

#5

Hi,

Your transformer is ok it is just the way you pass the data that fails.
The following does the trick:

data = {'input': {'img_ref': img_rgb,
                  'img_unalign': img_a,
                 }
             }

align_step = Step(name='Image Aligner',
                  transformer=ImageAlignmentTransformer(),
                  input_data=['input'],
                  cache_dirpath='.cache')

output = align_step.fit_transform(data)

Notice that you neet to run fit_transform(data) intead of transform unless your transformer is persisted in the ‘.cache/transformers’ directory. So the second time you run the same code it will load that trained transformer in.

In our case there is not much to fit so it’s just a dump/load for consistency with the rest of the framework.


#6

Great! Thank you! I finally get it now. Perhaps you can use this as an example as a tutorial for writing custom transformers for people who are new to using the Steps library. Wrapping existing sklearn transformers is one thing, but writing a simple unit test like this one, for custom transformers, will help as well.

What was not immediately obvious to me was that we’re actually calling the align_step.fit_transform passing the entire input dictionary. Now the structure and api for declaring steps makes sense.

Here is the completed unit test:

    def test_image_alignment_step(self):
    logger.info("Loading images")

    # read a sample
    sample_id = '6120_2_2'  # 6140_3_1, 6120_2_2, 6070_2_2

    # read rgb image
    img_rgb = tiff.imread('{}/three_band/{}.tif'.format(self.dataset_dir, sample_id)).transpose([1, 2, 0])
    nw, nh, _ = img_rgb.shape
    logger.info('rgb shape: {}'.format(img_rgb.shape))

    # read a-band image, low resolution
    img_a = tiff.imread("{}/sixteen_band/{}_A.tif".format(self.dataset_dir, sample_id)).transpose([1, 2, 0])
    logger.info('a-band image shape: {}'.format(img_a.shape))

    # upscale a-band image using opencv, note cv2.resize dsize is (nh, nw)
    rescaled_a = cv2.resize(img_a, (nh, nw), interpolation=cv2.INTER_CUBIC)
    logger.info('rescaled a-band image shape: {}'.format(rescaled_a.shape))

    data_train = {'input':
                    {
                        'img_ref': img_rgb,
                        'img_unalign': rescaled_a,
                    }
                 }

    align_step = Step(name='Image Aligner',
                      transformer=ImageAlignmentTransformer(),
                      input_data=['input'],
                      # adapter={'img_ref'     : ([('input', 'img_ref')]),
                      #          'img_unalign' : ([('input', 'img_unalign')])},
                      cache_dirpath='./cache')

    output = align_step.fit_transform(data_train)

    # print the output shape
    logger.info('img_align shape: {}'.format(output['img_align'].shape))

#7

Hi @jakub_czakon,

I have a quick question about arguments for a transformer’s constructor and fit methods when all you need is the transform method.

If I have the following image resize function,

def resize(img, w, h, interpolation=cv2.INTER_CUBIC):
    """
    Resize an image using OpenCV.
    Note that cv2.resize dsize is (h, w).
    :param img: Input image
    :param w: width
    :param h: height
    :param interpolation: cv2 interpolation type
    :return: resized image
    """
    return cv2.resize(img, (h, w), interpolation=interpolation)

Should the Step transformer have the parameters img, w, h, interpolation for the init(self) and fit(self) methods as well:

i.e. should it look like this without the img, w, h, interpolation for the init(self) and fit(self) methods:

class ImageResizeTransformer(BaseTransformer):
    def __init__(self):
        self.resize = resize

    def fit(self, **kwargs):
        return self

    def transform(self, img, w, h, interpolation):
        img_resize = self.resize(img=img, w=w, h=h, interpolation=interpolation)
        return {'img_resize': img_resize}

    def save(self, filepath):
        joblib.dump(self.resize, filepath)

    def load(self, filepath):
        self.resize = joblib.load(filepath)
        return self

or should I add the img, w, h, interpolation for the fit(self) methods, even if they are not used?