flickr_filter.py

import imghdr
import itertools
import os
import sys
from time import time
import csv
import PIL.Image as im
import numpy as np
import scipy.io
import scipy.misc as spm

sys.path.append('..')
sys.path.append('../slim')
sys.path.append('../object_detection')
# Notice: you need to clone TF-slim and Tensorflow Object Detection API
# into data_processing:
# https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/slim
# https://github.com/tensorflow/models/tree/master/research/object_detection

import cv2
import coco_data_provider as coco
import tensorflow as tf
from slim.nets import nets_factory
from slim.preprocessing import preprocessing_factory

from object_detection.utils import label_map_util


inception_ckpt_path = '../../inception_resnet_v2/inception_resnet_v2_2016_08_30.ckpt'

slim = tf.contrib.slim
tf.logging.set_verbosity(tf.logging.INFO)


def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape(
        (im_height, im_width, 3)).astype(np.uint8), im_width, im_height


def get_imagenet_class_labels():
    synset_list = [s.strip() for s in open('./imagenet_lsvrc_2015_synsets.txt', 'r').readlines()]
    num_synsets_in_ilsvrc = len(synset_list)
    assert num_synsets_in_ilsvrc == 1000

    synset_to_human_list = open('./imagenet_metadata.txt', 'r').readlines()
    num_synsets_in_all_imagenet = len(synset_to_human_list)
    assert num_synsets_in_all_imagenet == 21842

    synset_to_human = {}
    for s in synset_to_human_list:
        parts = s.strip().split('\t')
        assert len(parts) == 2
        synset = parts[0]
        human = parts[1]
        synset_to_human[synset] = human

    label_index = 1
    labels_to_names = {0: 'background'}
    for synset in synset_list:
        name = synset_to_human[synset]
        labels_to_names[label_index] = name
        label_index += 1

    return labels_to_names


def check_jpg_vadility_single(path):
    if imghdr.what(path) == 'jpg':
        return True
    return False


def check_jpg_vadility(path):
    file_list = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
    invalid_file_list = []
    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")
    for i in range(len(file_list)):
        if i % 5000 == 0:
            curr_time = time()
            elapsed = curr_time - prev_time
            print(
                "Now at iteration %d. Elapsed time: %.5fs." % (i, elapsed))
            prev_time = curr_time
            print(len(invalid_file_list))
        try:
            img = im.open(os.path.join(path, file_list[i]))
            format = img.format.lower()
            if format != 'jpg' and format != 'jpeg':
                raise ValueError
            # img.load()
        except:
            invalid_file_list.append(file_list[i])
    return invalid_file_list


def build_imagenet_graph(path):
    tf.reset_default_graph()
    print(path)

    filename_queue = tf.train.string_input_producer(tf.train.match_filenames_once(path + "/*.jpg"),
                                                    num_epochs=1, shuffle=False, capacity=100)
    image_reader = tf.WholeFileReader()
    image_file_name, image_file = image_reader.read(filename_queue)

    image = tf.image.decode_jpeg(image_file, channels=3, fancy_upscaling=True)

    model_name = 'inception_resnet_v2'
    network_fn = nets_factory.get_network_fn(model_name, is_training=False, num_classes=1001)

    preprocessing_name = model_name
    image_preprocessing_fn = preprocessing_factory.get_preprocessing(preprocessing_name, is_training=False)

    eval_image_size = network_fn.default_image_size

    image = image_preprocessing_fn(image, eval_image_size, eval_image_size)

    filenames, images = tf.train.batch([image_file_name, image], batch_size=100, num_threads=2, capacity=500)
    logits, _ = network_fn(images)

    variables_to_restore = slim.get_variables_to_restore()
    predictions = tf.argmax(logits, 1)

    return filenames, logits, predictions, variables_to_restore


def filter_by_imagenet(path, cls_name):
    labels_dict = get_imagenet_class_labels()
    filenames, logits, predictions, variables_to_restore = build_imagenet_graph(path)
    saver = tf.train.Saver(variables_to_restore)
    output_filename_list = []
    counter = 0

    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())

        saver.restore(sess, inception_ckpt_path)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        while True:
            try:
                filename_list, logit_array, prediction_list = sess.run([filenames, logits, predictions])
            except Exception as e:
                break

            if counter % 5000 == 0:
                print("Evaluated %d files" % counter)
                print(len(output_filename_list))

            prediction_dict = {os.path.split(filename)[1]: labels_dict[prediction] for filename, prediction in
                               zip(filename_list, prediction_list)}
            for i, j in prediction_dict.items():
                j = [p.strip() for p in j.lower().split(',')]
                if cls_name.lower() in j and len(j) == 1:
                    output_filename_list.append(i.decode('ascii'))

            counter += 100

        coord.request_stop()
        coord.join(threads)

        return output_filename_list


# SSD filter for COCO classes in Tensorflow instead of Caffe.
# Not fully functional yet. It will not output filtered filenames.
def filter_by_coco(path, cls_name):
    TEST_IMAGE_PATHS = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]

    counter = 0
    output_filename_list = []

    tf.reset_default_graph()
    print(path)

    PATH_TO_CKPT = '../ssd_inception_v2/frozen_inference_graph.pb'
    PATH_TO_LABELS = os.path.join('../../object_detection/data', 'mscoco_label_map.pbtxt')
    NUM_CLASSES = 90

    # Label map
    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
                                                                use_display_name=True)
    category_index = label_map_util.create_category_index(categories)
    for i in range(80):
        if categories[i]['name'] == cls_name:
            cls_index = categories[i]['id']

    # Load graph
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        # Input queue
        filename_queue = tf.train.string_input_producer(tf.train.match_filenames_once(path + "/*.jpg"),
                                                        num_epochs=1, shuffle=False, capacity=100)
        image_reader = tf.WholeFileReader()
        image_file_name, image_file = image_reader.read(filename_queue)

        image = tf.image.decode_jpeg(image_file, channels=3, fancy_upscaling=True)
        image0 = tf.image.resize_image_with_crop_or_pad(image, 500, 500)
        image = tf.image.resize_images(image0, [250, 250], method=tf.image.ResizeMethod.BILINEAR)
        image = tf.cast(image, tf.uint8)

        filenames, images = tf.train.batch([image_file_name, image], batch_size=20, num_threads=2, capacity=500)

        # Graph Def
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='', input_map={'image_tensor:0': images})

    # Time counter
    prev_time = float("-inf")
    curr_time = float("-inf")

    with detection_graph.as_default():
        with tf.Session(graph=detection_graph, config=config) as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())

            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess=sess, coord=coord)

            # Definite input and output Tensors for detection_graph
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            # Each box represents a part of the image where a particular object was detected.
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            # Each score represent how level of confidence for each of the objects.
            # Score is shown on the result image, together with the class label.
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
            for image_path in TEST_IMAGE_PATHS:
                if counter % 5 == 0:
                    curr_time = time()
                    elapsed = curr_time - prev_time
                    print(
                        "Now at iteration %d. Elapsed time: %.5fs." % (counter, elapsed))
                    prev_time = curr_time

                image = im.open(image_path)
                # the array based representation of the image will be used later in order to prepare the
                # result image with boxes and labels on it.
                image_np, im_width, im_height = load_image_into_numpy_array(image)
                image_np = scipy.misc.imresize(image_np, 0.5, 'bilinear')
                # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
                image_np_expanded = np.expand_dims(image_np, axis=0)
                # Actual detection.
                (boxes, scores, classes, num) = sess.run(
                    [detection_boxes, detection_scores, detection_classes, num_detections],
                    feed_dict={image_tensor: image_np_expanded})

                # Filter results
                boxes = np.squeeze(boxes)
                classes = np.squeeze(classes).astype(np.int32),
                scores = np.squeeze(scores)
                idx = np.logical_and(scores > 0.9, classes == cls_index)
                portion = np.prod(boxes[idx], axis=1) / (im_width * im_height)

                if portion.size > 0:
                    print()

                counter += 1


def filter_images(flickr_dir, cls_name):
    imagenet_classes = [i[:-1] for i in open('./imagenet_share_classes.txt').readlines()]
    coco_classes = [i[:-1] for i in open('./coco_share_classes.txt').readlines()]

    this_dir = os.path.join(flickr_dir, cls_name)

    invalid_file_list = []
    print("Invalid file number: %d" % len(invalid_file_list))
    for file_name in invalid_file_list:
        os.remove(os.path.join(this_dir, file_name))

    if cls_name in imagenet_classes:
        output_filename_list = filter_by_imagenet(this_dir, cls_name)
    elif cls_name in coco_classes:
        output_filename_list = filter_by_coco(this_dir, cls_name)
    else:
        raise NotImplementedError

    file_list = [f for f in os.listdir(this_dir) if os.path.isfile(os.path.join(this_dir, f))]
    print(len(file_list) - len(output_filename_list))
    for file_name in file_list:
        if file_name not in output_filename_list:
            os.remove(os.path.join(this_dir, file_name))


config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False,
                        intra_op_parallelism_threads=4)
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.9

with open('./imagenet_share_classes.txt', 'r') as f:
    classes_list = [i[:-1] for i in f.readlines()]

if __name__ == '__main__':
    # # Imagenet
    # labels_dict = get_imagenet_class_labels()
    # labels_list = [label + '\n' for i, label in labels_dict.items()]
    # with open('./imagenet_classes.txt', 'w') as f:
    #     f.writelines(labels_list)

    # COCO
    labels_list = [cls['name'] for cls in coco.get_all_images_data_categories(split='train')[2]]

    # with open('./all_classes', 'r') as f:
    #     sketchy_class_list = [i[:-1] for i in f.readlines()]
    #
    # filtered_classes = []
    # for cls in sketchy_class_list:
    #     for large_cls in labels_list:
    #         large_cls_names = [i.strip() for i in large_cls.split(',')]
    #         for name in large_cls_names:
    #             if cls.lower() == name.lower() and cls.lower() not in filtered_classes:
    #                 filtered_classes.append(cls + '\n')
    #
    # with open('./coco_share_classes.txt', 'w') as f:
    #     f.writelines(filtered_classes)

    # Inference
    filter_range = (8, 12)

    class_list = ['airplane']
    for class_name in class_list:
        filter_images('../flickr_output', class_name)