occlusion_heatmap.py

import matplotlib
matplotlib.use('Agg')
from skimage.io import imread
matplotlib.rcParams.update({'font.size': 2})
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid

import sys
import numpy as np
import scipy.ndimage as nd
import scipy.misc
from nyanbar import NyanBar

import pdb

HOME = '/home/ilyakavalerov/'
CAFFE_ROOT = HOME+'install_caffe/caffe/'
# sys.path.insert(0, CAFFE_ROOT + 'python')
import caffe

MINIBATCH_SIZE = 256
MODEL_FILE = '../val2.prototxt'
PRETRAINED = '../food_alexnet_train_iter_25000.caffemodel'
INPUT_IMAGE = HOME+'code/fundus/data/train/cent_crop_227/1000016.png'

FRY_IMAGE = HOME+'code/fundus/data/train/rand/1049122.png'

def get_resized_imagenet_mean(out_size=[227,227]):
    # returns resized shape: (d,d,3)
    IMAGE_MEAN = '../imagenet_mean.binaryproto'
    a=caffe.io.caffe_pb2.BlobProto()
    file=open(IMAGE_MEAN,'rb')
    data = file.read()
    a.ParseFromString(data)
    mean=a.data
    mean=np.asarray(mean)
    # Unsure: is this the shape of the mean?
    mean=mean.reshape(3,256,256) # shape of imagenet_mean.binaryproto
    # resize
    zoom = [1] + (np.array(out_size, dtype=float) / np.array([256,256], dtype=float)).tolist()
    resized_mean = nd.zoom(mean, zoom, order=1)
    return np.rollaxis(resized_mean, 0,3)

class OcclusionStudy(object):
    def __init__(self, image, occlusion_patch_size=21, image_shape_override=None):
        """
        image shape should be dxdx3
        """
        assert(occlusion_patch_size % 2 == 1) # for simplicity
        self.occlusion_patch_size = occlusion_patch_size
        self.image = image
        self.image_shape = image_shape_override or self.image.shape
        self.channel_means = self.image.mean(axis=(0,1))

        patched_img_pad = (occlusion_patch_size - 1) / 2
        patched_img_dim = self.image_shape[0] - patched_img_pad*2
        self.num_patch_centers = (patched_img_dim)**2 # number of images in test set

        self.patch_starts = [divmod(n, patched_img_dim) for n in xrange(self.num_patch_centers)]

    def nth_patch(self, n):
        i_start,j_start = self.patch_starts[n]
        i_end = i_start + self.occlusion_patch_size
        j_end = j_start + self.occlusion_patch_size

        return(i_start,i_end,j_start,j_end)

    def nth_occluded_image(self, n):
        i_start,i_end,j_start,j_end = self.nth_patch(n)
        img = self.image
        img[i_start:i_end, j_start:j_end, :] = self.channel_means
        return img

    def accumulate_patches_into_heatmaps(self, all_test_output, outpath):
        """
        assumes all_test_output is a vector of the class probability of
        interest at each of the patch locations
        """
        accumulator = np.zeros(self.image_shape[:2])
        normalizer = np.zeros(self.image_shape[:2])
        for n in xrange(self.num_patch_centers):
            i_start,i_end,j_start,j_end = self.nth_patch(n)

            accumulator[i_start:i_end, j_start:j_end] += all_test_output[n]
            normalizer[i_start:i_end, j_start:j_end] += 1
        normalized_img = accumulator / normalizer

        plt.imshow(normalized_img, interpolation="nearest", vmin=0, vmax=1)
        print("Saving figure as: %s" % outpath)
        plt.savefig(outpath, dpi=600, bbox_inches='tight')

def load_imagenet_network():
    """
    load example Imagenet network

    only setup needed is to be in caffe source root and to run:
    ./scripts/download_model_binary.py models/bvlc_reference_caffenet
    ./data/ilsvrc12/get_ilsvrc_aux.sh

    from: http://nbviewer.ipython.org/github/BVLC/caffe/blob/master/examples/00-classification.ipynb
    """
    net = caffe.Net(CAFFE_ROOT + 'models/bvlc_reference_caffenet/deploy.prototxt',
                CAFFE_ROOT + 'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel',
                caffe.TEST)
    # input preprocessing: 'data' is the name of the input blob == net.inputs[0]
    transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
    transformer.set_transpose('data', (2,0,1))
    transformer.set_mean('data', np.load(CAFFE_ROOT + 'python/caffe/imagenet/ilsvrc_2012_mean.npy').mean(1).mean(1)) # mean pixel
    transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
    transformer.set_channel_swap('data', (2,1,0))  # the reference model has channels in BGR order instead of RGB

    net.blobs['data'].reshape(MINIBATCH_SIZE,3,227,227)
    return(net, transformer)

    # 281 is the class for tabby, tabby cat
    # net.blobs['data'].data[...] = transformer.preprocess('data', caffe.io.load_image(CAFFE_ROOT + 'examples/images/cat.jpg'))
    # out = net.forward()
    # compare to labels
    # imagenet_labels_filename = CAFFE_ROOT + 'data/ilsvrc12/synset_words.txt'
    # labels = np.loadtxt(imagenet_labels_filename, str, delimiter='\t')
    # top_k = net.blobs['prob'].data[0].flatten().argsort()[-1:-6:-1]

def load_foodnet_network():
    # mean = get_resized_imagenet_mean()
    # net = caffe.Classifier(MODEL_FILE, PRETRAINED, mean=mean)
    # net = caffe.Net(MODEL_FILE, PRETRAINED, caffe.TEST)
    raise(NotImplementedError)

def save_all_occluded_images(input_image, outpath, occlusion_patch_size=41, every_ith_saved=6):
    """
    eg:
    input_image = imread(FRY_IMAGE)
    save_all_occluded_images(input_image, '/home/ilyakavalerov/from_swami/caffe/occlusions/fry2/', 41, 6)
    """
    os = OcclusionStudy(input_image, occlusion_patch_size)
    pb = NyanBar(tasks=os.num_patch_centers)
    for n in range(os.num_patch_centers):
        if n % every_ith_saved == 0:
            img_cpy = imread(FRY_IMAGE)
            i_start,i_end,j_start,j_end = os.nth_patch(n)
            img_cpy[i_start:i_end, j_start:j_end, :] = os.channel_means
            outfile = outpath + ('%i.png' % n)
            scipy.misc.toimage(img_cpy).save(outfile, "PNG")
        pb.task_done()
    pb.finish()

if __name__ == '__main__':
    net, transformer = load_imagenet_network()

    caffe.set_device(0)
    caffe.set_mode_gpu()

    input_image = imread(CAFFE_ROOT + 'examples/images/cat.jpg')
    class_of_interest = 281
    cropped_centered_image = transformer.preprocess('data', input_image)

    os = OcclusionStudy(np.rollaxis(cropped_centered_image,0,3), occlusion_patch_size=31)

    all_predictions = []
    macrobatches = int(os.num_patch_centers / float(MINIBATCH_SIZE)) + 1

    pb = NyanBar(tasks=macrobatches)
    for macrobatch in range(macrobatches):
        print "%i/%i" % (macrobatch,macrobatches)

        net.blobs['data'].data[...] = cropped_centered_image
        for i in range(MINIBATCH_SIZE):
            # to avoid index out of range error on last macrobatch
            n = ((macrobatch * MINIBATCH_SIZE) + i) % os.num_patch_centers

            i_start,i_end,j_start,j_end = os.nth_patch(n)

            net.blobs['data'].data[i, 0, i_start:i_end, j_start:j_end] = os.channel_means[0]
            net.blobs['data'].data[i, 1, i_start:i_end, j_start:j_end] = os.channel_means[1]
            net.blobs['data'].data[i, 2, i_start:i_end, j_start:j_end] = os.channel_means[2]

        out = net.forward()
        class_probabilities = out['prob'][:,class_of_interest]
        # last macrobatch check
        if macrobatch == macrobatches-1:
            leftovers = os.num_patch_centers - MINIBATCH_SIZE*(macrobatches-1)
            all_predictions.extend(class_probabilities[:leftovers])
        else:
            all_predictions.extend(class_probabilities)
        pb.task_done()
    pb.finish()

    # output result
    outpath = './class%i_heatmap.png' % class_of_interest
    all_predictions = np.array(all_predictions)
    normed_predictions = all_predictions / all_predictions.max()
    os.accumulate_patches_into_heatmaps(normed_predictions, outpath=outpath)

    pdb.set_trace()