First commit of flickr crawler and post processing scripts

Author: wchen342

data_processing/coco_data_provider.py

@@ -0,0 +1,119 @@
+import os
+import sys
+
+sys.path.append('.')
+sys.path.append('./pycocotools')
+import pycocotools.coco as coco
+
+pascal_classes = ['person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep', 'aeroplane',
+                  'bicycle', 'boat', 'bus', 'car', 'motorbike', 'train', 'bottle', 'chair',
+                  'dining table', 'potted plant', 'sofa', 'tv/monitor']
+pascal_classes_mapped = ['person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep', 'airplane',
+                         'bicycle', 'boat', 'bus', 'car (sedan)', 'motorcycle', 'train', 'bottle', 'chair',
+                         'dining table', 'potted plant', 'sofa', 'tv']
+
+images_dir = '../Datasets/COCO/coco-master/images'
+anno_dir = '../Datasets/COCO/coco-master/annotations'
+itrain_2014 = 'instances_train2014.json'
+ival_2014 = 'instances_val2014.json'
+
+
+def img_info_list_to_dict(input_list):
+    dic = {}
+    for i in input_list:
+        image_id = i['id']
+        assert image_id not in dic.keys()
+        dic[image_id] = i
+    return dic
+
+
+def get_all_images_data_categories(split, catIds=[]):
+    if split == 'train':
+        COCO = coco.COCO(annotation_file=os.path.join(anno_dir, itrain_2014))
+    elif split == 'test':
+        COCO = coco.COCO(annotation_file=os.path.join(anno_dir, ival_2014))
+    if len(catIds) == 0:
+        return COCO.loadImgs(COCO.getImgIds()), COCO.loadAnns(COCO.getAnnIds()), COCO.loadCats(COCO.getCatIds())
+    else:
+        return COCO.loadImgs(ids=COCO.getImgIds(catIds=catIds)), COCO.loadAnns(
+            ids=COCO.getAnnIds(catIds=catIds)), COCO.loadCats(COCO.getCatIds(catIds=catIds))
+
+
+def expand_bbox(bbox, max_height, max_width, frac):
+    assert len(bbox) == 4
+    half_width = round(bbox[2] / 2)
+    half_height = round(bbox[3] / 2)
+    mid_x = bbox[0] + half_width
+    mid_y = bbox[1] + half_height
+
+    x_min = max(0, mid_x - half_width * frac)
+    y_min = max(0, mid_y - half_height * frac)
+    x_max = min(max_width, mid_x + half_width * frac)
+    y_max = min(max_height, mid_y + half_height * frac)
+    return [round(x_min), round(y_min), round(x_max), round(y_max)]
+
+
+def get_shared_classes(input=None, print_out=False, output_file=True):
+    if input is None:
+        ret = get_all_images_data_categories('train')[2]
+    else:
+        ret = input
+    coco_classes = [cls['name'] for cls in ret]
+    class_dict = {item['name']: item for item in ret}
+    # Convert 'car' to 'car (sedan)' for comparison with Sketchy
+    coco_classes[coco_classes.index('car')] = 'car (sedan)'
+    with open('../../shared_classes', 'r') as f:
+        shared_classes = [cls[:-1].replace('_', ' ') for cls in f.readlines()]
+
+    if print_out:
+        print(len(shared_classes))
+        print(shared_classes)
+    if output_file:
+        shared_classes2 = [(cls + '\n') for cls in shared_classes if cls in coco_classes]
+        with open('../../shared_classes2', 'w') as f:
+            f.writelines(shared_classes2)
+
+    shared_classes = [cls for cls in shared_classes if cls in coco_classes]
+    print([cls for cls in shared_classes if cls in coco_classes and cls in pascal_classes_mapped])
+    shared_classes[shared_classes.index('car (sedan)')] = 'car'  # Convert 'car' back
+    output_dict = {class_dict[cls]['id']: class_dict[cls] for cls in shared_classes}
+    return shared_classes, output_dict
+
+
+def get_bbox():
+    img_info, seg_info, cat_info = get_all_images_data_categories('train')
+    img_info = img_info_list_to_dict(img_info)
+    shared_classes, cls_info = get_shared_classes(input=cat_info, print_out=False, output_file=False)
+    bbox_list = {cls: [] for cls in shared_classes}
+    for object in seg_info:
+        category_id = object['category_id']
+        if category_id not in cls_info:
+            continue
+
+        category_name = cls_info[category_id]['name']
+        bbox = object['bbox']
+        image_id = object['image_id']
+        iscrowd = object['iscrowd']
+
+        this_img_info = img_info[image_id]
+        file_name = this_img_info['file_name']
+        height = this_img_info['height']
+        width = this_img_info['width']
+        bbox = expand_bbox(bbox, height, width, 1.5)
+
+        bbox_list[category_name].append({
+            'image_id': image_id,
+            'category_name': category_name,
+            'category_id': category_id,
+            'iscrowd': iscrowd,
+            'file_name': file_name,
+            'height': height,
+            'width': width,
+            'bbox': bbox,
+        })
+
+    return bbox_list
+
+
+if __name__ == '__main__':
+    get_bbox()

data_processing/coco_share_classes.txt

@@ -0,0 +1,19 @@
+airplane
+apple
+bear
+bench
+bicycle
+cat
+car
+chair
+couch
+cow
+dog
+elephant
+giraffe
+horse
+knife
+motorcycle
+scissors
+sheep
+spoon

data_processing/edge_detection/PostprocessHED.m

@@ -0,0 +1,80 @@
+% Modified from pix2pix scripts: https://github.com/phillipi/pix2pix/blob/master/scripts/edges
+
+%%% Prerequisites
+% You need to get the cpp file edgesNmsMex.cpp from https://raw.githubusercontent.com/pdollar/edges/master/private/edgesNmsMex.cpp
+% and compile it in Matlab: mex edgesNmsMex.cpp
+% You also need to download and install Piotr's Computer Vision Matlab Toolbox:  https://pdollar.github.io/toolbox/
+
+%%% parameters
+% hed_mat_dir: the hed mat file directory (the output of 'batch_hed.py')
+% edge_dir: the output HED edges directory
+% image_width: resize the edge map to [image_width, image_width] 
+% threshold: threshold for image binarization (default 25.0/255.0)
+% small_edge: remove small edges (default 5)
+
+function [] = PostprocessHED(hed_mat_dir, edge_dir, image_width, threshold, small_edge)
+
+if ~exist(edge_dir, 'dir')
+    mkdir(edge_dir);
+end
+fileList = dir(fullfile(hed_mat_dir, '*.mat'));
+nFiles = numel(fileList);
+fprintf('find %d mat files\n', nFiles);
+
+parfor n = 1 : nFiles
+    if mod(n, 1000) == 0
+        fprintf('process %d/%d images\n', n, nFiles);
+    end
+    fileName = fileList(n).name;
+    filePath = fullfile(hed_mat_dir, fileName);
+    jpgName = strrep(fileName, '.mat', '.png');
+    edge_path = fullfile(edge_dir, jpgName);
+    
+    if ~exist(edge_path, 'file')
+        E = GetEdge(filePath);
+        E = imresize(E,[image_width,image_width]);
+        E_simple = SimpleEdge(E, threshold, small_edge);
+        E_simple = uint8(E_simple*255);
+        imwrite(E_simple, edge_path);
+    end
+end
+end
+
+
+
+
+function [E] = GetEdge(filePath)
+matdata = load(filePath);
+E = 1-matdata.predict;
+end
+
+function [E4] = SimpleEdge(E, threshold, small_edge)
+if nargin <= 1
+    threshold = 25.0/255.0;
+end
+
+if nargin <= 2
+    small_edge = 5;
+end
+
+if ndims(E) == 3
+    E = E(:,:,1);
+end
+
+E1 = 1 - E;
+E2 = EdgeNMS(E1);
+E3 = double(E2>=max(eps,threshold));
+E3 = bwmorph(bwmorph(E3,'thin',inf), 'clean');
+E4 = bwareaopen(E3, small_edge);
+E4 = bwmorph(bwmorph(E4,'shrink'), 'spur');
+E4=1-E4;
+end
+
+function [E_nms] = EdgeNMS( E )
+E=single(E);
+[Ox,Oy] = gradient2(convTri(E,4));
+[Oxx,~] = gradient2(Ox);
+[Oxy,Oyy] = gradient2(Oy);
+O = mod(atan(Oyy.*sign(-Oxy)./(Oxx+1e-5)),pi);
+E_nms = edgesNmsMex(E,O,1,5,1.01,1);
+end

data_processing/edge_detection/batch_hed.py

@@ -0,0 +1,122 @@
+# HED batch processing script
+# Modified from pix2pix scripts: https://github.com/phillipi/pix2pix/tree/master/scripts/edges
+# Notice that this script processes images in batches instead of one-by-one, which is slightly faster.
+# However, it will discard the last n images where 0 < n < batch_size.
+# Originally modified from https://github.com/s9xie/hed/blob/master/examples/hed/HED-tutorial.ipynb
+# Step 1: download the hed repo: https://github.com/s9xie/hed
+# Step 2: download the models and protoxt, and put them under {caffe_root}/examples/hed/
+# Step 3: put this script under {caffe_root}/examples/hed/
+# Step 4: run the following script: 
+#       python batch_hed.py --images_dir=/data/to/path/photos/ --hed_mat_dir=/data/to/path/hed_mat_files/
+# Step 5: run the MATLAB post-processing script "PostprocessHED.m"
+# The code sometimes crashes after computation is done. Error looks like "Check failed: ... driver shutting down". You can just kill the job. 
+# For large images, it will produce gpu memory issue. Therefore, you better resize the images before running this script. 
+
+import numpy as np
+import scipy.misc
+from PIL import Image
+import matplotlib.pyplot as plt
+import matplotlib.pylab as pylab
+import matplotlib.cm as cm
+import scipy.io
+import os
+import cv2
+import argparse
+from time import time
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='batch proccesing: photos->edges')
+    parser.add_argument('--caffe_root', dest='caffe_root', help='caffe root',
+                        default='../hed-master', type=str)
+    parser.add_argument('--caffemodel', dest='caffemodel', help='caffemodel',
+                        default='../hed-master/examples/hed/hed_pretrained_bsds.caffemodel',
+                        type=str)
+    parser.add_argument('--prototxt', dest='prototxt', help='caffe prototxt file',
+                        default='../hed-master/examples/hed/deploy.prototxt', type=str)
+    parser.add_argument('--images_dir', dest='images_dir',
+                        default='../flickr_coco/airplane',
+                        help='directory to store input photos', type=str)
+    parser.add_argument('--hed_mat_dir', dest='hed_mat_dir', default='../flickr_hed_coco/mat/airplane',
+                        help='directory to store output hed edges in mat file',
+                        type=str)
+    parser.add_argument('--border', dest='border', help='padding border', type=int, default=128)
+    parser.add_argument('--gpu_id', dest='gpu_id', help='gpu id', type=int, default=0)
+    args = parser.parse_args()
+    return args
+
+
+args = parse_args()
+for arg in vars(args):
+    print('[%s] =' % arg, getattr(args, arg))
+# Make sure that caffe is on the python path:  
+caffe_root = args.caffe_root  # this file is expected to be in {caffe_root}/examples/hed/
+import sys
+sys.path.insert(0, os.path.join(caffe_root, 'python'))
+
+# Suppress output
+os.environ['GLOG_minloglevel'] = '2'
+
+import caffe
+import scipy.io as sio
+
+if not os.path.exists(args.hed_mat_dir):
+    print('create output directory %s' % args.hed_mat_dir)
+    os.makedirs(args.hed_mat_dir)
+
+imgList = os.listdir(args.images_dir)
+nImgs = len(imgList)
+print('#images = %d' % nImgs)
+
+caffe.set_mode_gpu()
+caffe.set_device(args.gpu_id)
+# load net
+net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)
+# pad border
+border = args.border
+
+# Time counter
+prev_time = float("-inf")
+curr_time = float("-inf")
+image_list = []
+batch_size = 8
+
+for i in range(nImgs):
+    if i % 500 == 0:
+        print('processing image %d/%d' % (i, nImgs))
+        curr_time = time()
+        elapsed = curr_time - prev_time
+        print(
+            "Now at iteration %d. Elapsed time: %.5fs. Average time: %.5fs/iter" % (i, elapsed, elapsed / 500.))
+        prev_time = curr_time
+
+    if i % batch_size == 0 and i > 0:
+        in_ = np.concatenate(image_list, axis=0)
+
+        # shape for input (data blob is N x C x H x W), set data
+        net.blobs['data'].reshape(*in_.shape)
+        net.blobs['data'].data[...] = in_
+        # run net and take argmax for prediction
+        net.forward()
+        fuse = np.squeeze(net.blobs['sigmoid-fuse'].data)
+        # get rid of the border
+        fuse = fuse[:, border:-border, border:-border]
+
+        for j in range(batch_size):
+            # save hed file to the disk
+            name, ext = os.path.splitext(imgList[i - batch_size + j])
+            sio.savemat(os.path.join(args.hed_mat_dir, name + '.mat'), {'predict': fuse[j]})
+
+        image_list = []
+
+    im = cv2.imread(os.path.join(args.images_dir, imgList[i]), cv2.IMREAD_COLOR)
+    im = cv2.resize(im, (256, 256), interpolation=cv2.INTER_AREA)
+
+    in_ = im.astype(np.float32)
+    in_ = np.pad(in_, ((border, border), (border, border), (0, 0)), 'reflect')
+
+    in_ = in_[:, :, ::-1]
+    in_ -= np.array((104.00698793, 116.66876762, 122.67891434))
+    in_ = np.expand_dims(in_.transpose((2, 0, 1)), axis=0)
+
+    image_list.append(in_)