Clean up code

lib/__init__.py

@@ -0,0 +1 @@
+from .models import *

visualize.py

@@ -1,12 +1,13 @@
 import argparse
 import os
+import math
 import cupy
 import numpy as np
-from chainer import serializers, cuda
-from chainer import Variable
-from models.VGG import VGG
-from utils import imgutil
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
 from PIL import Image
+from chainer import Variable, serializers, cuda
+from lib.models import VGG
 
 
 def parse_args():
@@ -15,10 +16,46 @@ def parse_args():
     parser.add_argument('--out-dirname', type=str, default='results')
     parser.add_argument('--image-filename', type=str, default='images/cat.jpg')
     parser.add_argument('--model-filename', type=str, default='VGG.model')
-
     return parser.parse_args()
 
 
+def save_im(filename, im):
+    im = np.rollaxis(im, 0, 3)  # (c, h, w) -> (h, w, c)
+    im = Image.fromarray(im)
+    im.save(filename)
+
+
+def save_ims(filename, ims, dpi=100, scale=0.5):
+    n, c, h, w = ims.shape
+
+    rows = int(math.ceil(math.sqrt(n)))
+    cols = int(round(math.sqrt(n)))
+
+    fig, axes = plt.subplots(rows, cols, figsize=(w*cols/dpi*scale, h*rows/dpi*scale), dpi=dpi)
+
+    for i, ax in enumerate(axes.flat):
+        if i < n:
+            ax.imshow(ims[i].transpose((1, 2, 0)))
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.axis('off')
+
+    plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0.1, hspace=0.1)
+    plt.savefig(filename, dpi=dpi, bbox_inces='tight', transparent=True)
+    plt.clf()
+    plt.close()
+
+
+def tile_ims(filename, directory):
+
+    """Load all images in the given directory and tile them into one."""
+
+    ims = [mpimg.imread(os.path.join(directory, f)) for f in sorted(os.listdir(directory))]
+    ims = np.array(ims)
+    ims = ims.transpose((0, 3, 1, 2))  # (n, h, w, c) -> (n, c, h ,w)
+    save_ims(filename, ims)
+
+
 def read_im(filename):
 
     """Return a preprocessed (averaged and resized to VGG) sample image."""
@@ -34,35 +71,60 @@ def read_im(filename):
     return im
 
 
-def get_activations(model, x, layer_idx):
+def visualize_layer_activations(model, x, layer_idx):
 
     """Compute the activations for each feature map for the given layer for
     this particular image. Note that the input x should be a mini-batch
     of size one, i.e. a single image.
     """
 
-    if model._device_id >= 0:  # GPU
+    if model._device_id and model._device_id >= 0:  # Using GPU
         x = cupy.array(x)
 
     a = model.activations(Variable(x), layer_idx)
 
-    if model._device_id >= 0:
+    if model._device_id and model._device_id >= 0:
         a = cupy.asnumpy(a)
 
-    # Center at 0 with std 0.1
+    # Normalize and rescale to [0, 255]
     a -= a.mean()
     a /= (a.std() + 1e-5)
     a *= 0.1
+    a -= a.min()
+    a /= a.max()
+    a *= 255
 
-    # Clip to [0, 1]
-    a += 0.5
-    a = np.clip(a, 0, 1)
+    return a.astype(np.uint8)
 
-    # To RGB
-    a *= 255
-    a = a.astype('uint8')
 
-    return a
+def visualize(out_dirname, im, model):
+
+    def create_subdir(layer_idx):
+        dirname = os.path.join(out_dirname, 'conv{}/'.format(layer_idx+1))
+        dirname = os.path.dirname(dirname)
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        return dirname
+
+    def save_to_ims(activations, dirname):
+        filename_len = len(str(len(activations)))
+
+        # Save each feature map activation as its own image
+        for i, activation in enumerate(activations):
+            filename = '{num:0{width}}.png'.format(num=i, width=filename_len)
+            filename = os.path.join(dirname, filename)
+            save_im(filename, activation)
+
+        # Save an image of all feature map activations in this layer
+        tiled_filename = os.path.join(out_dirname, 'conv{}.png'.format(layer_idx+1))
+        tile_ims(tiled_filename, dirname)
+
+    # Visualize each of the 5 convolutional layers in VGG
+    for layer_idx in range(5):
+        print('Visualizing activations for conv{}...'.format(layer_idx+1))
+        activations = visualize_layer_activations(model, im.copy(), layer_idx)
+        dirname = create_subdir(layer_idx)
+        save_to_ims(activations, dirname)
 
 
 def main(args):
@@ -71,38 +133,16 @@ def main(args):
     model_filename = args.model_filename
     image_filename = args.image_filename
 
-    print('Loading VGG model... ')
-
     model = VGG()
     serializers.load_hdf5(model_filename, model)
 
     if gpu >= 0:
         cuda.get_device(gpu).use()
         model.to_gpu()
 
-    sample_im = read_im(image_filename)
-
-    # Visualize each of the 5 convolutional layers in VGG
-    for layer_idx in range(5):
-        # Create the target directory if it doesn't already exist
-        dst_dir = os.path.join(out_dirname, 'layer_{}/'.format(layer_idx))
-        dst_dir = os.path.dirname(dst_dir)
-        if not os.path.exists(dst_dir):
-            os.makedirs(dst_dir)
-
-        print('Computing activations for layer {}...'.format(layer_idx))
-        activations = get_activations(model, sample_im, layer_idx)
-
-        filename_len = len(str(len(activations)))
-        for i, activation in enumerate(activations):
-            im = np.rollaxis(activation, 0, 3)  # c, h, w -> h, w, c
-            filename = os.path.join(dst_dir,
-                                    '{num:0{width}}.jpg'.format(num=i, width=filename_len))
-
-            imgutil.save_im(filename, im)
+    im = read_im(image_filename)
 
-        tiled_filename = os.path.join(out_dirname, 'layer_{}.jpg'.format(layer_idx))
-        imgutil.tile_ims(tiled_filename, dst_dir)
+    visualize(out_dirname, im, model)
 
 
 if __name__ == '__main__':