sampling_h_class.py

#!/usr/bin/env python
'''
Anh Nguyen <anh.ng8@gmail.com>
2016
'''

import os, sys
os.environ['GLOG_minloglevel'] = '2'    # suprress Caffe verbose prints

import settings
sys.path.insert(0, settings.caffe_root)
import caffe

import numpy as np
from numpy.linalg import norm
import scipy.misc, scipy.io
import argparse 
import util
from sampler import Sampler

import matplotlib.pyplot as plt
from decimal import Decimal

if settings.gpu:
    caffe.set_mode_gpu() # sampling on GPU 

class ClassConditionalSampler(Sampler):

    def __init__ (self):
        # Load the list of class names
        with open(settings.synset_file, 'r') as synset_file:
            self.class_names = [ line.split(",")[0].split(" ", 1)[1].rstrip('\n') for line in synset_file.readlines()]

        # Hard-coded list of layers that has been tested 
        self.fc_layers = ["fc6", "fc7", "fc8", "fc9", "loss3/classifier", "fc1000", "prob"]
        self.conv_layers = ["conv1", "conv2", "conv3", "conv4", "conv5"]


    def forward_backward_from_x_to_condition(self, net, end, image, condition):
        '''
        Forward and backward passes through 'net', the condition model p(y|x), here an image classifier. 
        '''

        unit = condition['unit']
        xy = condition['xy']

        dst = net.blobs[end]

        acts = net.forward(data=image, end=end)
        one_hot = np.zeros_like(dst.data)

        # Get the activations
        if end in self.fc_layers:
            layer_acts = acts[end][0]
        elif end in self.conv_layers:
            layer_acts = acts[end][0, :, xy, xy]

        best_unit = layer_acts.argmax()     # highest probability unit

        # Compute the softmax probs by hand because it's handy in case we want to condition on hidden units as well
        exp_acts = np.exp(layer_acts - np.max(layer_acts))
        probs = exp_acts / (1e-10 + np.sum(exp_acts, keepdims=True))

        # The gradient of log of softmax, log(p(y|x)), reduces to:
        softmax_grad = 1 - probs.copy()

        obj_prob = probs.flat[unit]

        # Assign the gradient 
        if end in self.fc_layers:
            one_hot.flat[unit] = softmax_grad[unit]
        elif end in self.conv_layers:
            one_hot[:, unit, xy, xy] = softmax_grad[unit]
        else:
            raise Exception("Invalid layer type!")
        
        dst.diff[:] = one_hot

        # Backpropagate the gradient to the image layer 
        diffs = net.backward(start=end, diffs=['data'])
        g = diffs['data'].copy()

        dst.diff.fill(0.)   # reset objective after each step

        # Info to be printed out in the below 'print_progress' method
        info = {
            'best_unit': best_unit,
            'best_unit_prob': probs.flat[best_unit]
        }
        return g, obj_prob, info 
    
    def forward_backward_from_h_to_condition(self, net, end, h_code, condition):
        '''
        Forward and backward passes through 'net' (h_classifier), the condition model p(y|h)
        '''

        unit = condition['unit']

        dst = net.blobs[end]

        acts = net.forward(fc6=h_code, end=end)
        one_hot = np.zeros_like(dst.data)

        # Get the activations
        if end in self.fc_layers:
            layer_acts = acts[end][0]
        else:
            raise Exception("Invalid layer type!")

        best_unit = layer_acts.argmax()     # highest probability unit

        # Compute the softmax probs by hand because it's handy in case we want to condition on hidden units as well
        exp_acts = np.exp(layer_acts - np.max(layer_acts))
        probs = exp_acts / (1e-10 + np.sum(exp_acts, keepdims=True))

        # The gradient of log of softmax, log(p(y|h)), reduces to:
        softmax_grad = 1 - probs.copy()

        obj_prob = probs.flat[unit]

        # Assign the gradient 
        if end in self.fc_layers:
            one_hot.flat[unit] = softmax_grad[unit]
        else:
            raise Exception("Invalid layer type!")
        
        dst.diff[:] = one_hot

        # Backpropagate the gradient to the image layer 
        diffs = net.backward(start=end, diffs=['fc6'])
        g = diffs['fc6'].copy()

        dst.diff.fill(0.)   # reset objective after each step

        # Info to be printed out in the below 'print_progress' method
        info = {
            'best_unit': best_unit,
            'best_unit_prob': probs.flat[best_unit]
        }
        return g, obj_prob, info 


    def get_label(self, condition):
        unit = condition['unit']
        return self.class_names[unit]


    def print_progress(self, i, info, condition, prob, grad):
        print "step: %04d\t max: %4s [%.2f]\t obj: %4s [%.2f]\t norm: [%.2f]" % ( i, info['best_unit'], info['best_unit_prob'], condition['unit'], prob, norm(grad) )

def get_code(encoder, path, layer, mask=None):
    '''
    Push the given image through an encoder (here, AlexNet) to get a code.
    '''

    # set up the inputs for the net: 
    image_size = encoder.blobs['data'].shape[2:]    # (1, 3, 227, 227)
    images = np.zeros_like(encoder.blobs["data"].data, dtype='float32')

    in_image = scipy.misc.imread(path)
    in_image = scipy.misc.imresize(in_image, (image_size[0], image_size[1]))
    images[0] = np.transpose(in_image, (2, 0, 1))   # convert to (3, 227, 227) format

    data = images[:,::-1]   # convert from RGB to BGR

    # subtract the ImageNet mean
    image_mean = scipy.io.loadmat('misc/ilsvrc_2012_mean.mat')['image_mean'] # (256, 256, 3)
    topleft = util.compute_topleft(image_size, image_mean.shape[:2])
    image_mean = image_mean[topleft[0]:topleft[0]+image_size[0], topleft[1]:topleft[1]+image_size[1]]   # crop the image mean
    data -= np.expand_dims(np.transpose(image_mean, (2,0,1)), 0)    # mean is already BGR

    if mask is not None:
        data *= mask

    # initialize the encoder
    encoder = caffe.Net(settings.encoder_definition, settings.encoder_weights, caffe.TEST)

    # extract the features
    encoder.forward(data=data)
    features = encoder.blobs[layer].data.copy()

    return features, data

def main():

    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--units', metavar='units', type=str, help='an unit to visualize e.g. [0, 999]')
    parser.add_argument('--n_iters', metavar='iter', type=int, default=10, help='Number of sampling steps per each unit')
    parser.add_argument('--threshold', metavar='w', type=float, default=-1.0, nargs='?', help='The probability threshold to decide whether to keep an image')
    parser.add_argument('--save_every', metavar='save_iter', type=int, default=1, help='Save a sample every N iterations. 0 to disable saving')
    parser.add_argument('--reset_every', metavar='reset_iter', type=int, default=0, help='Reset the code every N iterations')
    parser.add_argument('--lr', metavar='lr', type=float, default=2.0, nargs='?', help='Learning rate')
    parser.add_argument('--lr_end', metavar='lr', type=float, default=-1.0, nargs='?', help='Ending Learning rate')
    parser.add_argument('--epsilon1', metavar='lr', type=float, default=1.0, nargs='?', help='Prior')
    parser.add_argument('--epsilon2', metavar='lr', type=float, default=1.0, nargs='?', help='Condition')
    parser.add_argument('--epsilon3', metavar='lr', type=float, default=1.0, nargs='?', help='Noise')
    parser.add_argument('--epsilon4', metavar='lr', type=float, default=0.0, nargs='?', help='Context')
    parser.add_argument('--seed', metavar='n', type=int, default=0, nargs='?', help='Random seed')
    parser.add_argument('--xy', metavar='n', type=int, default=0, nargs='?', help='Spatial position for conv units')
    parser.add_argument('--opt_layer', metavar='s', type=str, help='Layer at which we optimize a code')
    parser.add_argument('--act_layer', metavar='s', type=str, default="fc8", help='Layer at which we activate a neuron')
    parser.add_argument('--init_file', metavar='s', type=str, default="None", help='Init image')
    parser.add_argument('--write_labels', action='store_true', default=False, help='Write class labels to images')
    parser.add_argument('--output_dir', metavar='b', type=str, default=".", help='Output directory for saving results')
    parser.add_argument('--net_weights', metavar='b', type=str, default=settings.encoder_weights, help='Weights of the net being visualized')
    parser.add_argument('--net_definition', metavar='b', type=str, default=settings.encoder_definition, help='Definition of the net being visualized')

    args = parser.parse_args()

    # Default to constant learning rate
    if args.lr_end < 0:
        args.lr_end = args.lr

    # summary
    print "-------------"
    print " units: %s    xy: %s" % (args.units, args.xy)
    print " n_iters: %s" % args.n_iters
    print " reset_every: %s" % args.reset_every
    print " save_every: %s" % args.save_every
    print " threshold: %s" % args.threshold

    print " epsilon1: %s" % args.epsilon1
    print " epsilon2: %s" % args.epsilon2
    print " epsilon3: %s" % args.epsilon3
    print " epsilon4: %s" % args.epsilon4

    print " start learning rate: %s" % args.lr
    print " end learning rate: %s" % args.lr_end
    print " seed: %s" % args.seed
    print " opt_layer: %s" % args.opt_layer
    print " act_layer: %s" % args.act_layer
    print " init_file: %s" % args.init_file
    print "-------------"
    print " output dir: %s" % args.output_dir
    print " net weights: %s" % args.net_weights
    print " net definition: %s" % args.net_definition
    print "-------------"

    # encoder and generator for images 
    encoder = caffe.Net(settings.encoder_definition, settings.encoder_weights, caffe.TEST)
    generator = caffe.Net(settings.generator_definition, settings.generator_weights, caffe.TEST)

    # condition network, here an image classification net
    net = caffe.Classifier(args.net_definition, args.net_weights,
                             mean = np.float32([104.0, 117.0, 123.0]), # ImageNet mean
                             channel_swap = (2,1,0)) # the reference model has channels in BGR order instead of RGB
    h_net = caffe.Net("./nets/h_classifier/h_classifier.prototxt", "./nets/h_classifier/h_classifier.caffemodel", caffe.TEST)
    #h_net = caffe.Net("/home/damichoi/ml/ppgn/nets/h_classifier/h_classifier_Adam.prototxt", "/home/damichoi/ml/ppgn/nets/h_classifier/h_classifier_Adam.caffemodel", caffe.TEST)

    # Fix the seed
    np.random.seed(args.seed)

    # Sampler for class-conditional generation
    sampler = ClassConditionalSampler()
    inpainting = None

    if args.init_file != "None":

        # Pre-compute masks if we want to perform inpainting 
        if args.epsilon4 > 0:
            mask, neg = util.get_mask()
        else:
            neg = None

        # Get the code for the masked image
        start_code, start_image = get_code(encoder=encoder, path=args.init_file, layer=args.opt_layer, mask=neg)

        # Package settings for in-painting experiments
        if args.epsilon4 > 0:
            inpainting = {
                "mask"      : mask,
                "mask_neg"  : neg,
                "image"     : start_image,
                "epsilon4"  : args.epsilon4
            }

        print "Loaded init code: ", start_code.shape
    else:
        # shape of the code being optimized
        shape = generator.blobs[settings.generator_in_layer].data.shape
        start_code = np.random.normal(0, 1, shape)
        print ">>", np.min(start_code), np.max(start_code)

    # Separate the dash-separated list of units into numbers
    conditions = [ { "unit": int(u), "xy": args.xy } for u in args.units.split("_") ]     
    
    # Optimize a code via gradient ascent
#    output_image, list_samples = sampler.sampling( condition_net=net, image_encoder=encoder, image_generator=generator, 
#                        gen_in_layer=settings.generator_in_layer, gen_out_layer=settings.generator_out_layer, start_code=start_code, 
#                        n_iters=args.n_iters, lr=args.lr, lr_end=args.lr_end, threshold=args.threshold, 
#                        layer=args.act_layer, conditions=conditions,
#                        epsilon1=args.epsilon1, epsilon2=args.epsilon2, epsilon3=args.epsilon3,
#                        inpainting=inpainting,
#                        output_dir=args.output_dir, 
#                        reset_every=args.reset_every, save_every=args.save_every)
    
    output_image, list_samples, h, d_prior_mins, d_prior_maxs, d_condition_mins, d_condition_maxs, boundary_points  = sampler.h_sampling( condition_net=h_net, image_encoder=encoder, image_generator=generator, 
                        gen_in_layer=settings.generator_in_layer, gen_out_layer=settings.generator_out_layer, start_code=start_code, 
                        n_iters=args.n_iters, lr=args.lr, lr_end=args.lr_end, threshold=args.threshold, 
                        layer=args.act_layer, conditions=conditions,
                        epsilon1=args.epsilon1, epsilon2=args.epsilon2, epsilon3=args.epsilon3,
                        inpainting=inpainting,
                        output_dir=args.output_dir, 
                        reset_every=args.reset_every, save_every=args.save_every)
    
    ##################### send h through the net to verify class probability #####################
    image_shape = net.blobs['data'].data.shape
    generator_output_shape = generator.blobs[settings.generator_out_layer].data.shape
    image_size = util.get_image_size(image_shape)
    generator_output_size = util.get_image_size(generator_output_shape)
    topleft = util.compute_topleft(image_size, generator_output_size)
    x = output_image.copy()
    # Crop from 256x256 to 227x227
    cropped_x = x[:,:,topleft[0]:topleft[0]+image_size[0], topleft[1]:topleft[1]+image_size[1]]
    cropped_x_copy = cropped_x.copy()
    probs = net.forward(data=cropped_x_copy, end='prob')
    class_prob = probs['prob'][0][conditions[0]["unit"]]
    print("class is " + str(conditions[0]["unit"]))
    print("class probability is " + str(class_prob))
    ##############################################################################################
    
    #################### Plot gradients vs. num_iters ####################
    # plot the gradients
    plt.subplot(3, 1, 1)    #subplot(nrows, ncols, plot_number)
    x1 = np.linspace(0, args.n_iters, args.n_iters + 1, endpoint=True)
    plt.title('d_prior and d_condition')
    plt.plot(x1, d_prior_mins, color="blue", linewidth=2.0, linestyle="--", label='d_prior mins')
    plt.plot(x1, d_prior_maxs, color="blue", linewidth=2.0, linestyle="-", label='d_prior maxs')
    plt.plot(x1, d_condition_mins, color="red", linewidth=2.0, linestyle="--", label='d_condition mins')
    plt.plot(x1, d_condition_maxs, color="red", linewidth=2.0, linestyle="-", label='d_condition maxs')
    plt.legend()
    #plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=2, mode="expand", borderaxespad=0.)
    
    plt.subplot(3, 1, 2)
    x2 = np.linspace(0, args.n_iters, args.n_iters + 1, endpoint=True)
    plt.title('d_prior (scaled by eps1=' + '%.0e'%Decimal(args.epsilon1) + ') and d_condition (scaled by eps2=' + '%.0e'%Decimal(args.epsilon2) + ')')
    #plt.title('d_condition', fontsize=30)
    plt.plot(x2, d_condition_mins*args.epsilon2, color="red", linewidth=2.0, linestyle="--", label='d_condition mins')
    plt.plot(x2, d_condition_maxs*args.epsilon2, color="red", linewidth=2.0, linestyle="-", label='d_condition maxs')
    plt.plot(x2, d_prior_mins*args.epsilon1, color="blue", linewidth=2.0, linestyle="--", label='d_prior mins (scaled)')
    plt.plot(x2, d_prior_maxs*args.epsilon1, color="blue", linewidth=2.0, linestyle="-", label='d_prior maxs (scaled)')
    plt.legend()
            
    plt.subplot(3, 1, 3)
    x3 = np.linspace(14, args.n_iters, args.n_iters + 1 - 14, endpoint=True)
    plt.title('d_prior (scaled by eps1=' + '%.0e'%Decimal(args.epsilon1) + ') and d_condition (scaled by eps2=' + '%.0e'%Decimal(args.epsilon2) + ') from n_iter=14')
    #plt.title('d_condition from n_iter=14', fontsize=30)
    plt.plot(x3, d_condition_mins[14:]*args.epsilon2, color="red", linewidth=2.0, linestyle="--", label='d_condition mins')
    plt.plot(x3, d_condition_maxs[14:]*args.epsilon2, color="red", linewidth=2.0, linestyle="-", label='d_condition maxs')
    plt.plot(x3, d_prior_mins[14:]*args.epsilon1, color="blue", linewidth=2.0, linestyle="--", label='d_prior mins (scaled)')
    plt.plot(x3, d_prior_maxs[14:]*args.epsilon1, color="blue", linewidth=2.0, linestyle="-", label='d_prior maxs (scaled)')
    plt.xlabel('num iters')
    plt.legend()
    
    for i in xrange(args.n_iters):
        if i % 20 == 0:
            plt.annotate('(%s, %s)' %(i, d_condition_maxs[i]), xy=(i, d_condition_maxs[i] + 0.0005), textcoords='data')
            #plt.annotate('(%s, %s)' %(i, d_condition_mins[i]), xy=(i, d_condition_mins[i] - 0.0005), textcoords='data')

#    plt.title('% of boundary points')       
#    plt.plot(boundary_points/float(start_code.shape[1])*100)
#    plt.xlabel('num iters')
    plt.show()
    #plt.savefig("%s/gradients_plt.png")#, dpi=72)
    
    ####################################################################

    # Output image
    filename = "%s/%s_%04d_%04d_%s_h_%s_%s_%s_%s__%s.jpg" % (
            args.output_dir,
            args.act_layer, 
            conditions[0]["unit"],
            args.n_iters,
            args.lr,
            str(args.epsilon1),
            str(args.epsilon2),
            str(args.epsilon3),
            str(args.epsilon4),
            args.seed
        )

    if inpainting != None:
        output_image = util.stitch(start_image, output_image) 

    # Save the final image
    util.save_image(output_image, filename)
    print "%s/%s" % (os.getcwd(), filename)

    # Write labels to images
    print "Saving images..."
    for p in list_samples:
        img, name, label = p
        util.save_image(img, name)
        if args.write_labels:
            util.write_label_to_img(name, label)

if __name__ == '__main__':
    main()