data.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import os
import random

import numpy as np

import mxnet as mx
from mxnet import nd

def transform(data, target_wd, target_ht, is_train, box):
    """Crop and normnalize an image nd array."""
    if box is not None:
        x, y, w, h = box
        data = data[y:min(y+h, data.shape[0]), x:min(x+w, data.shape[1])]

    # Resize to target_wd * target_ht.
    data = mx.image.imresize(data, target_wd, target_ht)

    # Normalize in the same way as the pre-trained model.
    data = data.astype(np.float32) / 255.0
    data = (data - mx.nd.array([0.485, 0.456, 0.406])) / mx.nd.array([0.229, 0.224, 0.225])

    if is_train:
        if random.random() < 0.5:
            data = nd.flip(data, axis=1)
        data, _ = mx.image.random_crop(data, (224, 224))
    else:
        data, _ = mx.image.center_crop(data, (224, 224))

    # Transpose from (target_wd, target_ht, 3)
    # to (3, target_wd, target_ht).
    data = nd.transpose(data, (2, 0, 1))

    # If image is greyscale, repeat 3 times to get RGB image.
    if data.shape[0] == 1:
        data = nd.tile(data, (3, 1, 1))
    return data.reshape((1,) + data.shape)


class CUB200Iter(mx.io.DataIter):
    """Iterator for the CUB200-2011 dataset.
    Parameters
    ----------
    data_path : str,
        The path to dataset directory.
    batch_k : int,
        Number of images per class in a batch.
    batch_size : int,
        Batch size.
    batch_size : tupple,
        Data shape. E.g. (3, 224, 224).
    is_train : bool,
        Training data or testig data. Training batches are randomly sampled.
        Testing batches are loaded sequentially until reaching the end.
    """
    def __init__(self, data_path, batch_k, batch_size, data_shape, is_train):
        super(CUB200Iter, self).__init__(batch_size)
        self.data_shape = (batch_size,) + data_shape
        self.batch_size = batch_size
        self.provide_data = [('data', self.data_shape)]
        self.batch_k = batch_k
        self.is_train = is_train

        self.train_image_files = [[] for _ in range(100)]
        self.test_image_files = []
        self.test_labels = []
        self.boxes = {}
        self.test_count = 0

        with open(os.path.join(data_path, 'images.txt'), 'r') as f_img, \
             open(os.path.join(data_path, 'image_class_labels.txt'), 'r') as f_label, \
             open(os.path.join(data_path, 'bounding_boxes.txt'), 'r') as f_box:
            for line_img, line_label, line_box in zip(f_img, f_label, f_box):
                fname = os.path.join(data_path, 'images', line_img.strip().split()[-1])
                label = int(line_label.strip().split()[-1]) - 1
                box = [int(float(v)) for v in line_box.split()[-4:]]
                self.boxes[fname] = box

                # Following "Deep Metric Learning via Lifted Structured Feature Embedding" paper,
                # we use the first 100 classes for training, and the remaining for testing.
                if label < 100:
                    self.train_image_files[label].append(fname)
                else:
                    self.test_labels.append(label)
                    self.test_image_files.append(fname)

        self.n_test = len(self.test_image_files)

    def get_image(self, img, is_train):
        """Load and transform an image."""
        img_arr = mx.image.imread(img)
        img_arr = transform(img_arr, 256, 256, is_train, self.boxes[img])
        return img_arr

    def sample_train_batch(self):
        """Sample a training batch (data and label)."""
        batch = []
        labels = []
        num_groups = self.batch_size // self.batch_k

        # For CUB200, we use the first 100 classes for training.
        sampled_classes = np.random.choice(100, num_groups, replace=False)
        for i in range(num_groups):
            img_fnames = np.random.choice(self.train_image_files[sampled_classes[i]],
                                          self.batch_k, replace=False)
            batch += [self.get_image(img_fname, is_train=True) for img_fname in img_fnames]
            labels += [sampled_classes[i] for _ in range(self.batch_k)]

        return nd.concatenate(batch, axis=0), labels

    def get_test_batch(self):
        """Sample a testing batch (data and label)."""

        batch_size = self.batch_size
        batch = [self.get_image(self.test_image_files[(self.test_count*batch_size + i)
                                                      % len(self.test_image_files)],
                                is_train=False) for i in range(batch_size)]
        labels = [self.test_labels[(self.test_count*batch_size + i)
                                   % len(self.test_image_files)] for i in range(batch_size)]
        return nd.concatenate(batch, axis=0), labels

    def reset(self):
        """Reset an iterator."""
        self.test_count = 0

    def next(self):
        """Return a batch."""
        if self.is_train:
            data, labels = self.sample_train_batch()
        else:
            if self.test_count * self.batch_size < len(self.test_image_files):
                data, labels = self.get_test_batch()
                self.test_count += 1
            else:
                self.test_count = 0
                raise StopIteration
        return mx.io.DataBatch(data=[data], label=[labels])

def cub200_iterator(data_path, batch_k, batch_size, data_shape):
    """Return training and testing iterator for the CUB200-2011 dataset."""
    return (CUB200Iter(data_path, batch_k, batch_size, data_shape, is_train=True),
            CUB200Iter(data_path, batch_k, batch_size, data_shape, is_train=False))