extract_random_patches.py

from collections import namedtuple
from threading import Thread, Event
from queue import Queue, Empty
from typing import Optional
from os.path import exists, dirname
from os import makedirs
from io import BytesIO

import csv

import numpy as np
from skimage.transform import resize
from skimage.color import gray2rgb, rgb2gray
from skimage.filters import scharr
from sklearn.feature_extraction.image import extract_patches_2d
from PIL import Image

from cropping.edge_stats import get_edge_statistics
from cropping.shrink import crop
from cropping.normalize import normalize

import tensorflow as tf

from apiclient import discovery
from oauth2client.client import GoogleCredentials


OutputItem = namedtuple('OutputItem', ['image_id', 'patch_id', 'patch'])
ImageItem = namedtuple('ImageItem', ['image_id', 'bytes'])
FileItem = namedtuple('FileItem', ['image_id', 'bucket', 'object'])


class QueueWorker(Thread):
    def __init__(self, queue: Queue, stop: Event):
        self.__queue = queue  # type: Queue
        self.__stop = stop  # type: Event
        Thread.__init__(self)

    def run(self):
        while not (self.__stop.is_set() and self.__queue.empty()):
            try:
                item = self.__queue.get(block=True, timeout=1)
                self.process(item)
                self.__queue.task_done()
            except Empty:
                pass
        self.done()

    def process(self, item):
        pass

    def done(self):
        pass


class TFRecordWriteWorker(QueueWorker):
    def __init__(self, queue: Queue, stop: Event):
        self.__writer = None  # type: Optional[tf.python_io.TFRecordWriter]
        self.__writer_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
        self.__written_patches = 0  # type: int
        self.__max_written_patches = 100000  # 1000 patches are approximately 0.7 .. 1.1 MB
        self.__file_counter = -1

        QueueWorker.__init__(self, queue, stop)

    def __del__(self):
        self._close_writer()

    def process(self, item: OutputItem):
        feature = {
            'iid': _int64_feature(item.patch_id),
            'pid': _int64_feature(item.image_id),
            'raw': _bytes_feature(item.patch.tostring())
        }
        self._write_feature(feature)

    def done(self):
        self._close_writer()

    def _get_filename(self):
        filename = 'data/%05i.tfrecord.gz' % self.__file_counter
        d = dirname(filename)
        if not exists(d):
            makedirs(d)
        return filename

    def _write_feature(self, feature: tf.train.Feature):
        if (self.__written_patches >= self.__max_written_patches) or (self.__writer is None):
            self.__file_counter += 1
            self.__written_patches = 0
            self._close_writer()

            filename = self._get_filename()
            self.__writer = tf.python_io.TFRecordWriter(filename, options=self.__writer_options)

        example = tf.train.Example(features=tf.train.Features(feature=feature))
        self.__writer.write(example.SerializeToString())
        self.__written_patches += 1

    def _close_writer(self):
        if self.__writer is not None:
            self.__writer.close()
            self.__writer = None


class PatchExtractionWorker(QueueWorker):
    def __init__(self, input_queue: Queue, stop_input: Event, output_queue: Queue):
        self.__output_queue = output_queue  # type: Queue
        self.__patch_id = -1

        # patch sets
        PatchSet = namedtuple('PatchSet', ['size', 'count'])

        self.__target_size = 32
        self.__sets = [PatchSet(size=32, count=64),
                       PatchSet(size=64, count=64),
                       PatchSet(size=128, count=32)]

        QueueWorker.__init__(self, input_queue, stop_input)

    def process(self, item: ImageItem):

        with BytesIO(item.bytes) as buffer, Image.open(buffer) as pi:
            image = np.array(pi)

        # fetch one image
        image = gray2rgb(image)

        # apply energy cropping
        energy = scharr(rgb2gray(image))
        energy = normalize(energy)
        stats = get_edge_statistics(energy, edge_width=50)

        # crop the image
        cropped = crop(image, energy, threshold=None, stats=stats)
        height, width = image.shape[0], image.shape[1]

        # extract patches
        unique_patches = []
        extracted_patches = 0
        for patch_set in self.__sets:
            patch_size = patch_set.size

            # sanity check the patch size
            if (patch_size >= height) or (patch_size >= width):
                continue

            # extract the patches
            patch_count = patch_set.count
            try:
                patches = extract_patches_2d(cropped, (patch_size, patch_size), patch_count)
            except ValueError:
                continue
            extracted_patches += len(patches)

            # sort patches by variance
            patches = np.split(patches, patch_count, axis=0)
            patches = np.squeeze(patches)
            patches = sorted(patches, key=get_variance)

            # deduplication of exactly identical images might be something sane;
            # we might end up with half of the patches being identical due to border effects
            for patch in patches:
                if patch.shape[0] != self.__target_size:
                    patch = resize(patch, (self.__target_size, self.__target_size), preserve_range=True)
                    patch = np.floor(patch).astype(np.uint8)

                did_match = False
                for up in unique_patches:
                    if np.array_equal(patch, up):
                        did_match = True
                        break
                if not did_match:
                    unique_patches.append(patch)

        print('File %i: Added %i patches; skipped %i non-distinct patches.' % (item.image_id, len(unique_patches),
                                                                               extracted_patches - len(unique_patches)))

        # stash patches for serialization
        for patch in unique_patches:
            self.__patch_id += 1

            item = OutputItem(image_id=item.image_id, patch_id=self.__patch_id, patch=patch)
            self.__output_queue.put(item)

    def done(self):
        pass


class DownloadWorker(QueueWorker):
    def __init__(self, input_queue: Queue, stop_input: Event, output_queue: Queue, client):
        self.__output_queue = output_queue  # type: Queue
        self.__client = client
        QueueWorker.__init__(self, input_queue, stop_input)

    def process(self, item: FileItem):
        assert item is not None
        assert item.bucket is not None and item.object is not None
        assert item.image_id >= 0

        request = self.__client.objects().get_media(bucket=item.bucket, object=item.object)
        result = request.execute()
        item = ImageItem(image_id=item.image_id, bytes=result)
        self.__output_queue.put(item)

    def done(self):
        pass


def get_variance(patch: np.ndarray) -> float:
    v = np.var(patch)  # type: float
    return -v


def _int64_feature(value) -> tf.train.Feature:
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def _bytes_feature(value) -> tf.train.Feature:
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def main():
    # prepare the clients
    credentials = GoogleCredentials.from_stream('google-credentials.json')

    def create_client():
        return discovery.build('storage', 'v1', credentials=credentials)

    # prepare the worker
    image_id = -1

    store_queue = Queue(maxsize=2000)
    stop_output = Event()
    record_writer = TFRecordWriteWorker(store_queue, stop_output)

    image_queue = Queue(maxsize=4000)
    stop_images = Event()
    n_extractors = 32
    patch_extractors = [PatchExtractionWorker(input_queue=image_queue, stop_input=stop_images,
                                              output_queue=store_queue)
                        for _ in range(0, n_extractors)]

    file_queue = Queue(maxsize=400000)
    stop_files = Event()
    n_downloaders = 8
    download_workers = [DownloadWorker(input_queue=file_queue, stop_input=stop_files,
                                       output_queue=image_queue, client=create_client())
                        for _ in range(0, n_downloaders)]

    record_writer.start()
    for thread in patch_extractors:
        thread.start()
    for thread in download_workers:
        thread.start()

    # fetch all images
    bucket = 'product-pictures'
    client = create_client()
    request = client.objects().list(
        bucket=bucket,
        prefix='',
        fields='nextPageToken,items(name)')

    with open('data/index.tsv', 'w', newline='') as csv_file:
        writer = csv.writer(csv_file, delimiter='\t', quotechar='\"', quoting=csv.QUOTE_MINIMAL)

        object_page = -1
        while request is not None:
            object_page += 1
            print('Requesting page %i of objects (seen objects: %i)...' % (object_page,
                                                                           image_id if image_id >= 0 else 0))
            response = request.execute()
            request = client.objects().list_next(request, response)

            if 'items' not in response:
                continue

            for item in response['items']:
                image_id += 1
                image_file = item['name']
                writer.writerow([image_id, image_file])

                item = FileItem(bucket=bucket, object=image_file, image_id=image_id)
                file_queue.put(item)

            csv_file.flush()

    stop_files.set()
    for thread in download_workers:
        thread.join()

    stop_images.set()
    for thread in patch_extractors:
        thread.join()

    stop_output.set()
    record_writer.join()


if __name__ == '__main__':
    main()