Make this module importable; control over GPU options

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+*.ckpt
+*.pyc

README.md

@@ -1,4 +1,35 @@
-# Human Pose Estimation with TensorFlow
+# Pose Estimation with TensorFlow (PoseTF)
+
+## Preamble
+
+This is a fork of [@eldar's pose-tensorflow](https://github.com/eldar/pose-tensorflow) repository for tracking multiple body parts in images using the Human Body Pose Estimation algorithm. It also works great for animal pose (demonstrated in [Mathis et al, 2018](https://www.nature.com/articles/s41593-018-0209-y)).
+
+The purpose of this fork is to make the interface to the algorithm more modular and easier to integrate into a standard Python workflow.
+
+Here is an example Python script that shows how the new modular organization works.
+
+
+    # The entire module can now be imported
+    import PoseTF
+
+    # Train a network, specifying the configuration file to use,
+    # the amount of GPU memory to request, the GPU to use, and
+    # the number of snapshots to keep on disk.
+    PoseTF.training.train(
+        cfg_filename='pose_cfg.yaml', 
+        max_to_keep=None, # Keep all snapshots!
+        memfrac=.3,
+        disable_autotune=True,
+        choose_gpu=3,
+    )
+
+
+The script and the dataset and can be located in any directory. It no longer has to be located in the "models" subdirectory in this repository.
+
+The old syntax `python3 ../../../train.py` still works, but is no longer necessary.
+
+Original documentation is located below.
+
 
 ![](images/teaser.png)
 

__init__.py

@@ -0,0 +1,6 @@
+from . import dataset
+from . import nnet
+from . import util
+from . import config
+from . import training
+from . import test

config.py

@@ -5,7 +5,7 @@
 import yaml
 from easydict import EasyDict as edict
 
-import default_config
+from . import default_config
 
 
 cfg = default_config.cfg

dataset/__init__.py

@@ -0,0 +1,5 @@
+from . import factory
+from . import mpii
+#from . import mscoco # requires pycocotools
+from . import penn_action
+from . import pose_dataset

dataset/factory.py

@@ -1,4 +1,4 @@
-from dataset.pose_dataset import PoseDataset
+from .pose_dataset import PoseDataset
 
 
 def create(cfg):

dataset/mpii.py

@@ -1,4 +1,4 @@
-from dataset.pose_dataset import PoseDataset
+from .pose_dataset import PoseDataset
 
 
 class MPII(PoseDataset):

dataset/penn_action.py

@@ -1,7 +1,7 @@
 from functools import reduce
 import numpy as np
 
-from dataset.pose_dataset import PoseDataset, Batch
+from .pose_dataset import PoseDataset, Batch
 
 
 def merge_batch(batches):

dataset/pose_dataset.py

@@ -7,7 +7,8 @@
 from numpy import concatenate as cat
 
 import scipy.io as sio
-from scipy.misc import imread, imresize
+import imageio
+import scipy.ndimage
 
 
 class Batch(Enum):
@@ -248,7 +249,10 @@ def make_batch(self, data_item, scale, mirror):
         im_file = data_item.im_path
         logging.debug('image %s', im_file)
         logging.debug('mirror %r', mirror)
-        image = imread(im_file, mode='RGB')
+
+        # Read the image
+        # Return PIL.Image with shape (height, width, 3)
+        image = imageio.imread(im_file, pilmode='RGB')
 
         if self.has_gt:
             joints = np.copy(data_item.joints)
@@ -258,8 +262,24 @@ def make_batch(self, data_item, scale, mirror):
             image = image[crop[1]:crop[3] + 1, crop[0]:crop[2] + 1, :]
             if self.has_gt:
                 joints[:, 1:3] -= crop[0:2].astype(joints.dtype)
-
-        img = imresize(image, scale) if scale != 1 else image
+
+        # Resize the image
+        if scale == 1:
+            img = image
+        else:
+            # Zoom each color channel separately
+            img = np.array([
+                scipy.ndimage.zoom(image[:, :, channel], scale, mode='reflect')
+                for channel in range(image.shape[2])])
+            img = np.moveaxis(img, 0, -1)
+
+            #~ # Alternative method using PIL
+            #~ # PIL does (width, height) instead of (height, width)
+            #~ new_size = (
+                #~ int(np.rint(image.shape[1] * scale)),
+                #~ int(np.rint(image.shape[0] * scale)))
+            #~ img = np.array(PIL.Image.fromarray(image).resize(new_size))
+
         scaled_img_size = arr(img.shape[0:2])
 
         if mirror:
@@ -405,4 +425,4 @@ def compute_scmap_weights(self, scmap_shape, joint_id, data_item):
                     weights[:, :, j_id] = 1.0
         else:
             weights = np.ones(scmap_shape)
-        return weights
+        return weights

default_config.py

@@ -9,7 +9,7 @@
 cfg.weigh_only_present_joints = False
 cfg.mean_pixel = [123.68, 116.779, 103.939]
 cfg.shuffle = True
-cfg.snapshot_prefix = "snapshot"
+cfg.snapshot_prefix = "./snapshot"
 cfg.log_dir = "log"
 cfg.global_scale = 1.0
 cfg.location_refinement = False

nnet/__init__.py

@@ -0,0 +1,4 @@
+from . import losses
+from . import net_factory
+from . import pose_net
+from . import predict

nnet/net_factory.py

@@ -1,4 +1,4 @@
-from nnet.pose_net import PoseNet
+from .pose_net import PoseNet
 
 
 def pose_net(cfg):

nnet/pose_net.py

@@ -4,8 +4,8 @@
 import tensorflow.contrib.slim as slim
 from tensorflow.contrib.slim.nets import resnet_v1
 
-from dataset.pose_dataset import Batch
-from nnet import losses
+from ..dataset.pose_dataset import Batch
+from . import losses
 
 
 net_funcs = {'resnet_50': resnet_v1.resnet_v1_50,

nnet/predict.py

@@ -2,17 +2,26 @@
 
 import tensorflow as tf
 
-from nnet.net_factory import pose_net
+from .net_factory import pose_net
 
 
-def setup_pose_prediction(cfg):
+def setup_pose_prediction(cfg, memfrac=None):
     inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
 
     outputs = pose_net(cfg).test(inputs)
 
     restorer = tf.train.Saver()
 
-    sess = tf.Session()
+
+    # CR: ConfigProto for tensorflow session
+    tfcfg = tf.ConfigProto()
+    if memfrac is not None:
+        tfcfg.gpu_options.per_process_gpu_memory_fraction = memfrac
+
+    # CR: Initialize the tensorflow session
+    sess = tf.Session(config=tfcfg)
+
+    #sess = tf.Session()
 
     sess.run(tf.global_variables_initializer())
     sess.run(tf.local_variables_initializer())
@@ -78,4 +87,4 @@ def argmax_arrows_predict(scmap, offmat, pairwise_diff, stride):
                 pos_f8_end = (np.array(maxloc).astype('float') * stride + 0.5 * stride + difference)[::-1]
                 arrows[(joint_idx, joint_idx_end)] = (pos_f8, pos_f8_end)
 
-    return arrows
+    return arrows

pip-requirements.txt

@@ -0,0 +1,2 @@
+pyyaml
+easydict

requirements.txt

@@ -0,0 +1,7 @@
+scipy
+matplotlib
+cython
+scikit-image
+pyyaml
+tensorflow=1.13.1
+# install easydict with pip

test.py

@@ -6,11 +6,11 @@
 import scipy.io
 import scipy.ndimage
 
-from config import load_config
-from dataset.factory import create as create_dataset
-from dataset.pose_dataset import Batch
-from nnet.predict import setup_pose_prediction, extract_cnn_output, argmax_pose_predict
-from util import visualize
+from .config import load_config
+from .dataset.factory import create as create_dataset
+from .dataset.pose_dataset import Batch
+from .nnet.predict import setup_pose_prediction, extract_cnn_output, argmax_pose_predict
+from .util import visualize
 
 
 def test_net(visualise, cache_scoremaps):

train.py

@@ -1,140 +1,4 @@
-import logging
-import threading
-
-import tensorflow as tf
-import tensorflow.contrib.slim as slim
-
-from config import load_config
-from dataset.factory import create as create_dataset
-from nnet.net_factory import pose_net
-from nnet.pose_net import get_batch_spec
-from util.logging import setup_logging
-
-
-class LearningRate(object):
-    def __init__(self, cfg):
-        self.steps = cfg.multi_step
-        self.current_step = 0
-
-    def get_lr(self, iteration):
-        lr = self.steps[self.current_step][0]
-        if iteration == self.steps[self.current_step][1]:
-            self.current_step += 1
-
-        return lr
-
-
-def setup_preloading(batch_spec):
-    placeholders = {name: tf.placeholder(tf.float32, shape=spec) for (name, spec) in batch_spec.items()}
-    names = placeholders.keys()
-    placeholders_list = list(placeholders.values())
-
-    QUEUE_SIZE = 20
-
-    q = tf.FIFOQueue(QUEUE_SIZE, [tf.float32]*len(batch_spec))
-    enqueue_op = q.enqueue(placeholders_list)
-    batch_list = q.dequeue()
-
-    batch = {}
-    for idx, name in enumerate(names):
-        batch[name] = batch_list[idx]
-        batch[name].set_shape(batch_spec[name])
-    return batch, enqueue_op, placeholders
-
-
-def load_and_enqueue(sess, enqueue_op, coord, dataset, placeholders):
-    while not coord.should_stop():
-        batch_np = dataset.next_batch()
-        food = {pl: batch_np[name] for (name, pl) in placeholders.items()}
-        sess.run(enqueue_op, feed_dict=food)
-
-
-def start_preloading(sess, enqueue_op, dataset, placeholders):
-    coord = tf.train.Coordinator()
-
-    t = threading.Thread(target=load_and_enqueue,
-                         args=(sess, enqueue_op, coord, dataset, placeholders))
-    t.start()
-
-    return coord, t
-
-
-def get_optimizer(loss_op, cfg):
-    learning_rate = tf.placeholder(tf.float32, shape=[])
-
-    if cfg.optimizer == "sgd":
-        optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9)
-    elif cfg.optimizer == "adam":
-        optimizer = tf.train.AdamOptimizer(cfg.adam_lr)
-    else:
-        raise ValueError('unknown optimizer {}'.format(cfg.optimizer))
-    train_op = slim.learning.create_train_op(loss_op, optimizer)
-
-    return learning_rate, train_op
-
-
-def train():
-    setup_logging()
-
-    cfg = load_config()
-    dataset = create_dataset(cfg)
-
-    batch_spec = get_batch_spec(cfg)
-    batch, enqueue_op, placeholders = setup_preloading(batch_spec)
-
-    losses = pose_net(cfg).train(batch)
-    total_loss = losses['total_loss']
-
-    for k, t in losses.items():
-        tf.summary.scalar(k, t)
-    merged_summaries = tf.summary.merge_all()
-
-    variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
-    restorer = tf.train.Saver(variables_to_restore)
-    saver = tf.train.Saver(max_to_keep=5)
-
-    sess = tf.Session()
-
-    coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders)
-
-    train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph)
-
-    learning_rate, train_op = get_optimizer(total_loss, cfg)
-
-    sess.run(tf.global_variables_initializer())
-    sess.run(tf.local_variables_initializer())
-
-    # Restore variables from disk.
-    restorer.restore(sess, cfg.init_weights)
-
-    max_iter = int(cfg.multi_step[-1][1])
-
-    display_iters = cfg.display_iters
-    cum_loss = 0.0
-    lr_gen = LearningRate(cfg)
-
-    for it in range(max_iter+1):
-        current_lr = lr_gen.get_lr(it)
-        [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries],
-                                          feed_dict={learning_rate: current_lr})
-        cum_loss += loss_val
-        train_writer.add_summary(summary, it)
-
-        if it % display_iters == 0:
-            average_loss = cum_loss / display_iters
-            cum_loss = 0.0
-            logging.info("iteration: {} loss: {} lr: {}"
-                         .format(it, "{0:.4f}".format(average_loss), current_lr))
-
-        # Save snapshot
-        if (it % cfg.save_iters == 0 and it != 0) or it == max_iter:
-            model_name = cfg.snapshot_prefix
-            saver.save(sess, model_name, global_step=it)
-
-    sess.close()
-    coord.request_stop()
-    coord.join([thread])
-
+import PoseTF
 
 if __name__ == '__main__':
-    train()
+    PoseTF.training.train()

training/__init__.py

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