sampling_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

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", "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 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

    # 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
    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()