generate_patch_db.py

"""This script creates a leveldb database for a given UBC patch dataset. For
each patch we generate a key-value pair:
    key: the patch id (zero-starting line index in the info.txt file).
  value: a Caffe Datum containing the image patch and the metadata.

It will complain if the specified db already exists.

Example:
  python generate_patch_db.py data/phototour/liberty/info.txt \
    data/phototour/liberty/interest.txt \
    data/phototour/liberty data/leveldb/liberty.leveldb

"""

import leveldb, numpy as np, skimage
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from caffe.proto import caffe_pb2
from caffe.io import *


def ParseArgs():
    """Parse input arguments.
    """
    parser = ArgumentParser(description=__doc__,
                            formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('info_file',
                        help='Path to info.txt file in the dataset.')
    parser.add_argument('interest_file',
                        help='Path to interest.txt file in the dataset.')
    parser.add_argument('container_dir',
                        help='Patch to the directory of .bmp files.')
    parser.add_argument('output_db', help='Path to output database.')
    args = parser.parse_args()
    return args


def GetPatchImage(patch_id, container_dir):
    """Returns a 64 x 64 patch with the given patch_id. Catch container images to
       reduce loading from disk.
    """
    # Define constants. Each container image is of size 1024x1024. It packs at
    # most 16 rows and 16 columns of 64x64 patches, arranged from left to right,
    # top to bottom.
    PATCHES_PER_IMAGE = 16 * 16
    PATCHES_PER_ROW = 16
    PATCH_SIZE = 64

    # Calculate the container index, the row and column index for the given
    # patch.
    container_idx, container_offset = divmod(patch_id, PATCHES_PER_IMAGE)
    row_idx, col_idx = divmod(container_offset, PATCHES_PER_ROW)

    # Read the container image if it is not cached.
    if GetPatchImage.cached_container_idx != container_idx:
        GetPatchImage.cached_container_idx = container_idx
        GetPatchImage.cached_container_img = \
            skimage.img_as_ubyte(skimage.io.imread('%s/patches%04d.bmp' % \
                (container_dir, container_idx), as_grey=True))

    # Extract the patch from the image and return.
    patch_image = GetPatchImage.cached_container_img[ \
        PATCH_SIZE * row_idx:PATCH_SIZE * (row_idx + 1), \
        PATCH_SIZE * col_idx:PATCH_SIZE * (col_idx + 1)]
    return patch_image

# Static variables initialization for GetPatchImage.
GetPatchImage.cached_container_idx = None
GetPatchImage.cached_container_img = None


def main():
    # Parse input arguments.
    args = ParseArgs()

    # Read the 3Dpoint IDs from the info file.
    with open(args.info_file) as f:
        point_id = [int(line.split()[0]) for line in f]

    # Read the interest point from the interest file. The fields in each line
    # are: image_id, x, y, orientation, and scale. We parse all of them as float
    # even though image_id is integer.
    with open(args.interest_file) as f:
        interest = [[float(x) for x in line.split()] for line in f]

    # Create the output database, fail if exists.
    db = leveldb.LevelDB(args.output_db,
                         create_if_missing=True,
                         error_if_exists=True)

    # Add patches to the database in batch.
    batch = leveldb.WriteBatch()
    total = len(interest)
    processed = 0
    for i, metadata in enumerate(interest):
        datum = caffe_pb2.Datum()
        datum.channels, datum.height, datum.width = (1, 64, 64)

        # Extract the patch
        datum.data = GetPatchImage(i, args.container_dir).tostring()

        # Write 3D point ID into the label field.
        datum.label = point_id[i]

        # Write other metadata into float_data fields.
        datum.float_data.extend(metadata)
        batch.Put(str(i), datum.SerializeToString())
        processed += 1
        if processed % 1000 == 0:
            print processed, '/', total

            # Write the current batch.
            db.Write(batch, sync=True)

            # Verify the last written record.
            d = caffe_pb2.Datum()
            d.ParseFromString(db.Get(str(processed - 1)))
            assert (d.data == datum.data)

            # Start a new batch
            batch = leveldb.WriteBatch()
    db.Write(batch, sync=True)


if __name__ == '__main__':
    main()