update cpm

examples/ConvolutionalPoseMachines/README.md

+# Convolutional Pose Machines
+
+A script to load and run pretrained CPM model released by Shih-En. The original code in caffe is [here](https://github.com/shihenw/convolutional-pose-machines-release).
+
+## Usage:
+```
+# download the relased caffe model:
+wget http://pearl.vasc.ri.cmu.edu/caffe_model_github/model/_trained_MPI/pose_iter_320000.caffemodel
+wget https://github.com/shihenw/convolutional-pose-machines-release/raw/master/model/_trained_MPI/pose_deploy_resize.prototxt
+# convert the model to a dict:
+python -m tensorpack.utils.loadcaffe pose_deploy_resize.prototxt pose_iter_320000.caffemodel CPM-original.npy
+# run it on a test image, and produce output.jpg
+python load-cpm.py --load CPM-original.npy --input test.jpg
+```

examples/ConvolutionalPoseMachines/load-cpm.py

+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+# File: load-cpm.py
+# Author: Yuxin Wu <ppwwyyxxc@gmail.com>
+
+import cv2
+import tqdm
+import tensorflow as tf
+import numpy as np
+import os, argparse
+import matplotlib.pyplot as plt
+
+from tensorpack import *
+from tensorpack.utils.argtools import memoized
+from tensorpack.tfutils.symbolic_functions import *
+from tensorpack.tfutils.summary import *
+from tensorpack.dataflow.dataset import ILSVRCMeta
+
+_CM = plt.get_cmap('jet')
+def colorize(img, heatmap):
+    """ img: bgr, [0,255]
+        heatmap: [0,1]
+    """
+    heatmap = _CM(heatmap)[:,:,[2,1,0]] * 255.0
+    return img * 0.5 + heatmap * 0.5
+
+@memoized
+def get_gaussian_map():
+    sigma = 21
+    gaussian_map = np.zeros((368, 368), dtype='float32')
+    for x_p in range(368):
+        for y_p in range(368):
+            dist_sq = (x_p - 368/2) * (x_p - 368/2) + \
+                      (y_p - 368/2) * (y_p - 368/2)
+            exponent = dist_sq / 2.0 / (21**2)
+            gaussian_map[y_p, x_p] = np.exp(-exponent)
+    return gaussian_map.reshape((1,368,368,1))
+
+class Model(ModelDesc):
+    def _get_input_vars(self):
+        return [InputVar(tf.float32, (None, 368, 368, 3), 'input'),
+                InputVar(tf.float32, (None, 368, 368, 15), 'label'),
+                ]
+
+    def _build_graph(self, inputs):
+        image, label = inputs
+        image = image / 256.0 - 0.5
+
+        gmap = tf.constant(get_gaussian_map())
+        gmap = tf.pad(gmap, [[0,0],[0,1],[0,1],[0,0]])
+        pool_center = AvgPooling('mappool', gmap, 9, stride=8, padding='VALID')
+        with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu,
+                W_init=tf.random_normal_initializer(stddev=0.01)):
+            shared = (LinearWrap(image)
+                .Conv2D('conv1_1', 64)
+                .Conv2D('conv1_2', 64)
+                .MaxPooling('pool1', 2)
+                # 184
+                .Conv2D('conv2_1', 128)
+                .Conv2D('conv2_2', 128)
+                .MaxPooling('pool2', 2)
+                # 92
+                .Conv2D('conv3_1', 256)
+                .Conv2D('conv3_2', 256)
+                .Conv2D('conv3_3', 256)
+                .Conv2D('conv3_4', 256)
+                .MaxPooling('pool3', 2)
+                # 46
+                .Conv2D('conv4_1', 512)
+                .Conv2D('conv4_2', 512)
+                .Conv2D('conv4_3_CPM', 256)
+                .Conv2D('conv4_4_CPM', 256)
+                .Conv2D('conv4_5_CPM', 256)
+                .Conv2D('conv4_6_CPM', 256)
+                .Conv2D('conv4_7_CPM', 128)())
+
+        def add_stage(stage, l):
+            l = tf.concat(3, [l, shared, pool_center], name='concat_stage{}'.format(stage))
+            for i in range(1, 6):
+                l = Conv2D('Mconv{}_stage{}'.format(i, stage), l, 128)
+            l = Conv2D('Mconv6_stage{}'.format(stage), l, 128, kernel_shape=1)
+            l = Conv2D('Mconv7_stage{}'.format(stage),
+                    l, 15, kernel_shape=1, nl=tf.identity)
+            return l
+
+        with argscope(Conv2D, kernel_shape=7, nl=tf.nn.relu):
+            out1 = (LinearWrap(shared)
+                  .Conv2D('conv5_1_CPM', 512, kernel_shape=1)
+                  .Conv2D('conv5_2_CPM', 15, kernel_shape=1, nl=tf.identity)())
+            out2 = add_stage(2, out1)
+            out3 = add_stage(3, out2)
+            out4 = add_stage(4, out3)
+            out5 = add_stage(5, out4)
+            out6 = add_stage(6, out4)
+            resized_map = tf.image.resize_bilinear(out6,
+                    [368,368], name='resized_map')
+
+def run_test(model_path, img_file):
+    param_dict = np.load(model_path, encoding='latin1').item()
+    predict_func = OfflinePredictor(PredictConfig(
+        model=Model(),
+        session_init=ParamRestore(param_dict),
+        input_names=['input'],
+        output_names=['resized_map']
+    ))
+
+    im = cv2.imread(img_file, cv2.IMREAD_COLOR).astype('float32')
+    out = predict_func([[im]])[0][0]
+    hm = out[:,:,:14].sum(axis=2)
+    viz = colorize(im, hm)
+    cv2.imwrite("output.jpg", viz)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--load', required=True, help='.npy model file')
+    parser.add_argument('--input', required=True, help='input image')
+    args = parser.parse_args()
+    run_test(args.load, args.input)

tensorpack/utils/loadcaffe.py

@@ -123,9 +123,9 @@ def get_caffe_pb():
 if __name__ == '__main__':
     import argparse
     parser = argparse.ArgumentParser()
-    parser.add_argument('model')
-    parser.add_argument('weights')
-    parser.add_argument('output')
+    parser.add_argument('model', help='.prototxt file')
+    parser.add_argument('weights', help='.caffemodel file')
+    parser.add_argument('output', help='output npy file')
     args = parser.parse_args()
     ret = load_caffe(args.model, args.weights)