ADD file via upload

posegait/src/common

@@ -0,0 +1,135 @@
+from enum import Enum
+
+import tensorflow as tf
+import cv2
+
+
+regularizer_conv = 0.004
+regularizer_dsconv = 0.0004
+batchnorm_fused = True
+activation_fn = tf.nn.relu
+
+
+class CocoPart(Enum):
+    Nose = 0
+    Neck = 1
+    RShoulder = 2
+    RElbow = 3
+    RWrist = 4
+    LShoulder = 5
+    LElbow = 6
+    LWrist = 7
+    RHip = 8
+    RKnee = 9
+    RAnkle = 10
+    LHip = 11
+    LKnee = 12
+    LAnkle = 13
+    REye = 14
+    LEye = 15
+    REar = 16
+    LEar = 17
+    Background = 18
+
+
+class MPIIPart(Enum):
+    RAnkle = 0
+    RKnee = 1
+    RHip = 2
+    LHip = 3
+    LKnee = 4
+    LAnkle = 5
+    RWrist = 6
+    RElbow = 7
+    RShoulder = 8
+    LShoulder = 9
+    LElbow = 10
+    LWrist = 11
+    Neck = 12
+    Head = 13
+
+    @staticmethod
+    def from_coco(human):
+        # t = {
+        #     MPIIPart.RAnkle: CocoPart.RAnkle,
+        #     MPIIPart.RKnee: CocoPart.RKnee,
+        #     MPIIPart.RHip: CocoPart.RHip,
+        #     MPIIPart.LHip: CocoPart.LHip,
+        #     MPIIPart.LKnee: CocoPart.LKnee,
+        #     MPIIPart.LAnkle: CocoPart.LAnkle,
+        #     MPIIPart.RWrist: CocoPart.RWrist,
+        #     MPIIPart.RElbow: CocoPart.RElbow,
+        #     MPIIPart.RShoulder: CocoPart.RShoulder,
+        #     MPIIPart.LShoulder: CocoPart.LShoulder,
+        #     MPIIPart.LElbow: CocoPart.LElbow,
+        #     MPIIPart.LWrist: CocoPart.LWrist,
+        #     MPIIPart.Neck: CocoPart.Neck,
+        #     MPIIPart.Nose: CocoPart.Nose,
+        # }
+
+        t = [
+            (MPIIPart.Head, CocoPart.Nose),
+            (MPIIPart.Neck, CocoPart.Neck),
+            (MPIIPart.RShoulder, CocoPart.RShoulder),
+            (MPIIPart.RElbow, CocoPart.RElbow),
+            (MPIIPart.RWrist, CocoPart.RWrist),
+            (MPIIPart.LShoulder, CocoPart.LShoulder),
+            (MPIIPart.LElbow, CocoPart.LElbow),
+            (MPIIPart.LWrist, CocoPart.LWrist),
+            (MPIIPart.RHip, CocoPart.RHip),
+            (MPIIPart.RKnee, CocoPart.RKnee),
+            (MPIIPart.RAnkle, CocoPart.RAnkle),
+            (MPIIPart.LHip, CocoPart.LHip),
+            (MPIIPart.LKnee, CocoPart.LKnee),
+            (MPIIPart.LAnkle, CocoPart.LAnkle),
+        ]
+
+        pose_2d_mpii = []
+        visibilty = []
+        for mpi, coco in t:
+            if coco.value not in human.body_parts.keys():
+                pose_2d_mpii.append((0, 0))
+                visibilty.append(False)
+                continue
+            pose_2d_mpii.append((human.body_parts[coco.value].x, human.body_parts[coco.value].y))
+            visibilty.append(True)
+        return pose_2d_mpii, visibilty
+
+CocoPairs = [
+    (1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10), (1, 11),
+    (11, 12), (12, 13), (1, 0), (0, 14), (14, 16), (0, 15), (15, 17), (2, 16), (5, 17)
+]   # = 19
+CocoPairsRender = CocoPairs[:-2]
+CocoPairsNetwork = [
+    (12, 13), (20, 21), (14, 15), (16, 17), (22, 23), (24, 25), (0, 1), (2, 3), (4, 5),
+    (6, 7), (8, 9), (10, 11), (28, 29), (30, 31), (34, 35), (32, 33), (36, 37), (18, 19), (26, 27)
+ ]  # = 19
+
+CocoColors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
+              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
+              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
+
+
+def read_imgfile(path, width, height):
+    val_image = cv2.imread(path, cv2.IMREAD_COLOR)
+    if width is not None and height is not None:
+        val_image = cv2.resize(val_image, (width, height))
+    return val_image
+
+
+def get_sample_images(w, h):
+    val_image = [
+        read_imgfile('./images/p1.jpg', w, h),
+        read_imgfile('./images/p2.jpg', w, h),
+        read_imgfile('./images/p3.jpg', w, h),
+        read_imgfile('./images/golf.jpg', w, h),
+        read_imgfile('./images/hand1.jpg', w, h),
+        read_imgfile('./images/hand2.jpg', w, h),
+        read_imgfile('./images/apink1_crop.jpg', w, h),
+        read_imgfile('./images/ski.jpg', w, h),
+        read_imgfile('./images/apink2.jpg', w, h),
+        read_imgfile('./images/apink3.jpg', w, h),
+        read_imgfile('./images/handsup1.jpg', w, h),
+        read_imgfile('./images/p3_dance.png', w, h),
+    ]
+    return val_image

posegait/src/datum_pb2.py

@@ -0,0 +1,109 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: datum.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='datum.proto',
+  package='',
+  serialized_pb=_b('\n\x0b\x64\x61tum.proto\"\x81\x01\n\x05\x44\x61tum\x12\x10\n\x08\x63hannels\x18\x01 \x01(\x05\x12\x0e\n\x06height\x18\x02 \x01(\x05\x12\r\n\x05width\x18\x03 \x01(\x05\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\r\n\x05label\x18\x05 \x01(\x05\x12\x12\n\nfloat_data\x18\x06 \x03(\x02\x12\x16\n\x07\x65ncoded\x18\x07 \x01(\x08:\x05\x66\x61lse')
+)
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+
+
+
+_DATUM = _descriptor.Descriptor(
+  name='Datum',
+  full_name='Datum',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='channels', full_name='Datum.channels', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='height', full_name='Datum.height', index=1,
+      number=2, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='width', full_name='Datum.width', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='data', full_name='Datum.data', index=3,
+      number=4, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='label', full_name='Datum.label', index=4,
+      number=5, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='float_data', full_name='Datum.float_data', index=5,
+      number=6, type=2, cpp_type=6, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='encoded', full_name='Datum.encoded', index=6,
+      number=7, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=16,
+  serialized_end=145,
+)
+
+DESCRIPTOR.message_types_by_name['Datum'] = _DATUM
+
+Datum = _reflection.GeneratedProtocolMessageType('Datum', (_message.Message,), dict(
+  DESCRIPTOR = _DATUM,
+  __module__ = 'datum_pb2'
+  # @@protoc_insertion_point(class_scope:Datum)
+  ))
+_sym_db.RegisterMessage(Datum)
+
+
+# @@protoc_insertion_point(module_scope)

posegait/src/estimator.py

@@ -0,0 +1,465 @@
+import itertools
+import logging
+import math
+from collections import namedtuple
+
+import cv2
+import numpy as np
+import tensorflow as tf
+from scipy.ndimage import maximum_filter, gaussian_filter
+
+import common
+from common import CocoPairsNetwork, CocoPairs, CocoPart
+
+logger = logging.getLogger('TfPoseEstimator')
+logger.setLevel(logging.INFO)
+ch = logging.StreamHandler()
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+
+class Human:
+    """
+    body_parts: list of BodyPart
+    """
+    __slots__ = ('body_parts', 'pairs', 'uidx_list')
+
+    def __init__(self, pairs):
+        self.pairs = []
+        self.uidx_list = set()
+        self.body_parts = {}
+        for pair in pairs:
+            self.add_pair(pair)
+
+    @staticmethod
+    def _get_uidx(part_idx, idx):
+        return '%d-%d' % (part_idx, idx)
+
+    def add_pair(self, pair):
+        self.pairs.append(pair)
+        self.body_parts[pair.part_idx1] = BodyPart(Human._get_uidx(pair.part_idx1, pair.idx1),
+                                                   pair.part_idx1,
+                                                   pair.coord1[0], pair.coord1[1], pair.score)
+        self.body_parts[pair.part_idx2] = BodyPart(Human._get_uidx(pair.part_idx2, pair.idx2),
+                                                   pair.part_idx2,
+                                                   pair.coord2[0], pair.coord2[1], pair.score)
+        self.uidx_list.add(Human._get_uidx(pair.part_idx1, pair.idx1))
+        self.uidx_list.add(Human._get_uidx(pair.part_idx2, pair.idx2))
+
+    def is_connected(self, other):
+        return len(self.uidx_list & other.uidx_list) > 0
+
+    def merge(self, other):
+        for pair in other.pairs:
+            self.add_pair(pair)
+
+    def part_count(self):
+        return len(self.body_parts.keys())
+
+    def get_max_score(self):
+        return max([x.score for _, x in self.body_parts.items()])
+
+    def __str__(self):
+        return ' '.join([str(x) for x in self.body_parts.values()])
+
+
+class BodyPart:
+    """
+    part_idx : part index(eg. 0 for nose)
+    x, y: coordinate of body part
+    score : confidence score
+    """
+    __slots__ = ('uidx', 'part_idx', 'x', 'y', 'score')
+
+    def __init__(self, uidx, part_idx, x, y, score):
+        self.uidx = uidx
+        self.part_idx = part_idx
+        self.x, self.y = x, y
+        self.score = score
+
+    def get_part_name(self):
+        return CocoPart(self.part_idx)
+
+    def __str__(self):
+        return 'BodyPart:%d-(%.2f, %.2f) score=%.2f' % (self.part_idx, self.x, self.y, self.score)
+
+
+class PoseEstimator:
+    heatmap_supress = False
+    heatmap_gaussian = False
+    adaptive_threshold = False
+
+    NMS_Threshold = 0.15
+    Local_PAF_Threshold = 0.2
+    PAF_Count_Threshold = 5
+    Part_Count_Threshold = 4
+    Part_Score_Threshold = 4.5
+
+    PartPair = namedtuple('PartPair', [
+        'score',
+        'part_idx1', 'part_idx2',
+        'idx1', 'idx2',
+        'coord1', 'coord2',
+        'score1', 'score2'
+    ], verbose=False)
+    
+    #print('PartPair', PartPair)
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def non_max_suppression(plain, window_size=3, threshold=NMS_Threshold):
+        under_threshold_indices = plain < threshold
+        plain[under_threshold_indices] = 0
+        return plain * (plain == maximum_filter(plain, footprint=np.ones((window_size, window_size))))
+
+    @staticmethod
+    def estimate(heat_mat, paf_mat):
+        #print('heat_mat', heat_mat)
+        if heat_mat.shape[2] == 19:
+            heat_mat = np.rollaxis(heat_mat, 2, 0)
+        if paf_mat.shape[2] == 38:
+            paf_mat = np.rollaxis(paf_mat, 2, 0)
+
+        if PoseEstimator.heatmap_supress:
+            heat_mat = heat_mat - heat_mat.min(axis=1).min(axis=1).reshape(19, 1, 1)
+            heat_mat = heat_mat - heat_mat.min(axis=2).reshape(19, heat_mat.shape[1], 1)
+
+        if PoseEstimator.heatmap_gaussian:
+            heat_mat = gaussian_filter(heat_mat, sigma=0.5)
+
+        if PoseEstimator.adaptive_threshold:
+            _NMS_Threshold = max(np.average(heat_mat) * 4.0, PoseEstimator.NMS_Threshold)
+            _NMS_Threshold = min(_NMS_Threshold, 0.3)
+        else:
+            _NMS_Threshold = PoseEstimator.NMS_Threshold
+
+        # extract interesting coordinates using NMS.
+        coords = []     # [[coords in plane1], [....], ...]
+        for plain in heat_mat[:-1]:
+            nms = PoseEstimator.non_max_suppression(plain, 5, _NMS_Threshold)
+            coords.append(np.where(nms >= _NMS_Threshold))
+        #print('coords', coords)
+
+        # score pairs
+        pairs_by_conn = list()
+        for (part_idx1, part_idx2), (paf_x_idx, paf_y_idx) in zip(CocoPairs, CocoPairsNetwork):
+            pairs = PoseEstimator.score_pairs(
+                part_idx1, part_idx2,
+                coords[part_idx1], coords[part_idx2],
+                paf_mat[paf_x_idx], paf_mat[paf_y_idx],
+                heatmap=heat_mat,                
+                rescale=(1.0 / heat_mat.shape[2], 1.0 / heat_mat.shape[1])                
+            )
+            #print('pairs', pairs)
+
+            pairs_by_conn.extend(pairs)
+            #print('pairs by conn', pairs_by_conn)
+
+        # merge pairs to human
+        # pairs_by_conn is sorted by CocoPairs(part importance) and Score between Parts.
+        humans = [Human([pair]) for pair in pairs_by_conn]
+        #print('humans', humans)
+        while True:
+            merge_items = None
+            for k1, k2 in itertools.combinations(humans, 2):
+                if k1 == k2:
+                    continue
+                if k1.is_connected(k2):
+                    merge_items = (k1, k2)
+                    break
+
+            if merge_items is not None:
+                merge_items[0].merge(merge_items[1])
+                humans.remove(merge_items[1])
+            else:
+                break
+
+        # reject by subset count
+        humans = [human for human in humans if human.part_count() >= PoseEstimator.PAF_Count_Threshold]
+        #print('humans1', humans)
+        
+        # reject by subset max score
+        humans = [human for human in humans if human.get_max_score() >= PoseEstimator.Part_Score_Threshold]
+        #print('humans2', humans)
+
+        return humans
+
+    @staticmethod
+    def score_pairs(part_idx1, part_idx2, coord_list1, coord_list2, paf_mat_x, paf_mat_y, heatmap, rescale=(1.0, 1.0)):
+        connection_temp = []
+
+        cnt = 0
+        for idx1, (y1, x1) in enumerate(zip(coord_list1[0], coord_list1[1])):
+            for idx2, (y2, x2) in enumerate(zip(coord_list2[0], coord_list2[1])):
+                score, count = PoseEstimator.get_score(x1, y1, x2, y2, paf_mat_x, paf_mat_y)
+                cnt += 1
+                if count < PoseEstimator.PAF_Count_Threshold or score <= 0.0:
+                    continue
+                connection_temp.append(PoseEstimator.PartPair(
+                    score=score,
+                    part_idx1=part_idx1, part_idx2=part_idx2,
+                    idx1=idx1, idx2=idx2,
+                    coord1=(x1 * rescale[0], y1 * rescale[1]),
+                    coord2=(x2 * rescale[0], y2 * rescale[1]),
+                    score1=heatmap[part_idx1][y1][x1],
+                    score2=heatmap[part_idx2][y2][x2],
+                ))
+                #print('connection_temp', connection_temp)
+
+        connection = []
+        used_idx1, used_idx2 = set(), set()
+        for candidate in sorted(connection_temp, key=lambda x: x.score, reverse=True):
+            # check not connected
+            if candidate.idx1 in used_idx1 or candidate.idx2 in used_idx2:
+                continue
+            connection.append(candidate)
+            used_idx1.add(candidate.idx1)
+            used_idx2.add(candidate.idx2)
+            #print('connection', connection)
+
+        return connection
+
+    @staticmethod
+    def get_score(x1, y1, x2, y2, paf_mat_x, paf_mat_y):
+        __num_inter = 10
+        __num_inter_f = float(__num_inter)
+        dx, dy = x2 - x1, y2 - y1
+        normVec = math.sqrt(dx ** 2 + dy ** 2)
+
+        if normVec < 1e-4:
+            return 0.0, 0
+
+        vx, vy = dx / normVec, dy / normVec
+
+        xs = np.arange(x1, x2, dx / __num_inter_f) if x1 != x2 else np.full((__num_inter,), x1)
+        ys = np.arange(y1, y2, dy / __num_inter_f) if y1 != y2 else np.full((__num_inter,), y1)
+        xs = (xs + 0.5).astype(np.int8)
+        ys = (ys + 0.5).astype(np.int8)
+
+        # without vectorization
+        pafXs = np.zeros(__num_inter)
+        pafYs = np.zeros(__num_inter)
+        for idx, (mx, my) in enumerate(zip(xs, ys)):
+            pafXs[idx] = paf_mat_x[my][mx]
+            pafYs[idx] = paf_mat_y[my][mx]
+
+        # vectorization slow?
+        # pafXs = pafMatX[ys, xs]
+        # pafYs = pafMatY[ys, xs]
+
+        local_scores = pafXs * vx + pafYs * vy
+        thidxs = local_scores > PoseEstimator.Local_PAF_Threshold
+
+        return sum(local_scores * thidxs), sum(thidxs)
+
+
+class TfPoseEstimator:
+    ENSEMBLE = 'addup'        # average, addup
+
+    def __init__(self, graph_path, target_size=(320, 240)):
+        self.target_size = target_size
+
+        # load graph
+        with tf.gfile.GFile(graph_path, 'rb') as f:
+            graph_def = tf.GraphDef()
+            graph_def.ParseFromString(f.read())
+
+        self.graph = tf.get_default_graph()
+        tf.import_graph_def(graph_def, name='TfPoseEstimator')
+        self.persistent_sess = tf.Session(graph=self.graph)
+
+        # for op in self.graph.get_operations():
+        #     print(op.name)
+
+        self.tensor_image = self.graph.get_tensor_by_name('TfPoseEstimator/image:0')
+        self.tensor_output = self.graph.get_tensor_by_name('TfPoseEstimator/Openpose/concat_stage7:0')
+
+        self.heatMat = self.pafMat = None
+        
+        print('Inside TFPoseestimator')
+        
+        print('graph_path', graph_path)        
+        
+        # warm-up
+        self.persistent_sess.run(
+            self.tensor_output,
+            feed_dict={
+                self.tensor_image: [np.ndarray(shape=(target_size[1], target_size[0], 3), dtype=np.float32)]
+            }
+        )
+
+    def __del__(self):
+        self.persistent_sess.close()
+
+    @staticmethod
+    def _quantize_img(npimg):
+        npimg_q = npimg + 1.0
+        npimg_q /= (2.0 / 2**8)
+        # npimg_q += 0.5
+        npimg_q = npimg_q.astype(np.uint8)
+        return npimg_q
+
+    @staticmethod
+    def draw_humans(npimg, humans, imgcopy=False):
+        if imgcopy:
+            npimg = np.copy(npimg)
+        image_h, image_w = npimg.shape[:2]
+        centers = {}
+        for human in humans:
+            # draw point
+            for i in range(common.CocoPart.Background.value):
+                if i not in human.body_parts.keys():
+                    continue
+
+                body_part = human.body_parts[i]
+                center = (int(body_part.x * image_w + 0.5), int(body_part.y * image_h + 0.5))
+                centers[i] = center
+                cv2.circle(npimg, center, 3, common.CocoColors[i], thickness=3, lineType=8, shift=0)
+
+            # draw line
+            for pair_order, pair in enumerate(common.CocoPairsRender):
+                if pair[0] not in human.body_parts.keys() or pair[1] not in human.body_parts.keys():
+                    continue
+
+                npimg = cv2.line(npimg, centers[pair[0]], centers[pair[1]], common.CocoColors[pair_order], 3)
+
+        return npimg
+
+    def _get_scaled_img(self, npimg, scale):
+        get_base_scale = lambda s, w, h: max(self.target_size[0] / float(w), self.target_size[1] / float(h)) * s
+        img_h, img_w = npimg.shape[:2]
+
+        if scale is None:
+            if npimg.shape[:2] != (self.target_size[1], self.target_size[0]):
+                # resize
+                npimg = cv2.resize(npimg, self.target_size)
+            return [npimg], [(0.0, 0.0, 1.0, 1.0)]
+        elif isinstance(scale, float):
+            # scaling with center crop
+            base_scale = get_base_scale(scale, img_w, img_h)
+            npimg = cv2.resize(npimg, dsize=None, fx=base_scale, fy=base_scale)
+            ratio_x = (1. - self.target_size[0] / float(npimg.shape[1])) / 2.0
+            ratio_y = (1. - self.target_size[1] / float(npimg.shape[0])) / 2.0
+            roi = self._crop_roi(npimg, ratio_x, ratio_y)
+            return [roi], [(ratio_x, ratio_y, 1.-ratio_x*2, 1.-ratio_y*2)]
+        elif isinstance(scale, tuple) and len(scale) == 2:
+            # scaling with sliding window : (scale, step)
+            base_scale = get_base_scale(scale[0], img_w, img_h)
+            base_scale_w = self.target_size[0] / (img_w * base_scale)
+            base_scale_h = self.target_size[1] / (img_h * base_scale)
+            npimg = cv2.resize(npimg, dsize=None, fx=base_scale, fy=base_scale)
+            window_step = scale[1]
+            rois = []
+            infos = []
+            for ratio_x, ratio_y in itertools.product(np.arange(0., 1.01 - base_scale_w, window_step),
+                                                      np.arange(0., 1.01 - base_scale_h, window_step)):
+                roi = self._crop_roi(npimg, ratio_x, ratio_y)
+                rois.append(roi)
+                infos.append((ratio_x, ratio_y, base_scale_w, base_scale_h))
+            return rois, infos
+        elif isinstance(scale, tuple) and len(scale) == 3:
+            # scaling with ROI : (want_x, want_y, scale_ratio)
+            base_scale = get_base_scale(scale[2], img_w, img_h)
+            npimg = cv2.resize(npimg, dsize=None, fx=base_scale, fy=base_scale)
+            ratio_w = self.target_size[0] / float(npimg.shape[1])
+            ratio_h = self.target_size[1] / float(npimg.shape[0])
+
+            want_x, want_y = scale[:2]
+            ratio_x = want_x - ratio_w / 2.
+            ratio_y = want_y - ratio_h / 2.
+            ratio_x = max(ratio_x, 0.0)
+            ratio_y = max(ratio_y, 0.0)
+            if ratio_x + ratio_w > 1.0:
+                ratio_x = 1. - ratio_w
+            if ratio_y + ratio_h > 1.0:
+                ratio_y = 1. - ratio_h
+
+            roi = self._crop_roi(npimg, ratio_x, ratio_y)
+            return [roi], [(ratio_x, ratio_y, ratio_w, ratio_h)]
+
+    def _crop_roi(self, npimg, ratio_x, ratio_y):
+        target_w, target_h = self.target_size
+        h, w = npimg.shape[:2]
+        x = max(int(w*ratio_x-.5), 0)
+        y = max(int(h*ratio_y-.5), 0)
+        cropped = npimg[y:y+target_h, x:x+target_w]
+
+        cropped_h, cropped_w = cropped.shape[:2]
+        if cropped_w < target_w or cropped_h < target_h:
+            npblank = np.zeros((self.target_size[1], self.target_size[0], 3), dtype=np.uint8)
+
+            copy_x, copy_y = (target_w - cropped_w) // 2, (target_h - cropped_h) // 2
+            npblank[copy_y:copy_y+cropped_h, copy_x:copy_x+cropped_w] = cropped
+        else:
+            return cropped
+
+    def inference(self, npimg, scales=None):
+        if npimg is None:
+            raise Exception('The image is not valid. Please check your image exists.')
+
+        if not isinstance(scales, list):
+            scales = [None]
+
+        if self.tensor_image.dtype == tf.quint8:
+            # quantize input image
+            npimg = TfPoseEstimator._quantize_img(npimg)
+            pass
+
+        rois = []
+        infos = []
+        for scale in scales:
+            roi, info = self._get_scaled_img(npimg, scale)
+            # for dubug...
+            # print(roi[0].shape)
+            # cv2.imshow('a', roi[0])
+            # cv2.waitKey()
+            rois.extend(roi)
+            infos.extend(info)
+
+        logger.debug('inference+')
+        output = self.persistent_sess.run(self.tensor_output, feed_dict={self.tensor_image: rois})
+        heatMats = output[:, :, :, :19]
+        pafMats = output[:, :, :, 19:]
+        logger.debug('inference-')
+
+        output_h, output_w = output.shape[1:3]
+        max_ratio_w = max_ratio_h = 10000.0
+        for info in infos:
+            max_ratio_w = min(max_ratio_w, info[2])
+            max_ratio_h = min(max_ratio_h, info[3])
+        mat_w, mat_h = int(output_w/max_ratio_w), int(output_h/max_ratio_h)
+        resized_heatMat = np.zeros((mat_h, mat_w, 19), dtype=np.float32)
+        resized_pafMat = np.zeros((mat_h, mat_w, 38), dtype=np.float32)
+        resized_cntMat = np.zeros((mat_h, mat_w, 1), dtype=np.float32)
+        resized_cntMat += 1e-12
+
+        for heatMat, pafMat, info in zip(heatMats, pafMats, infos):
+            w, h = int(info[2]*mat_w), int(info[3]*mat_h)
+            heatMat = cv2.resize(heatMat, (w, h))
+            pafMat = cv2.resize(pafMat, (w, h))
+            x, y = int(info[0] * mat_w), int(info[1] * mat_h)
+
+            if TfPoseEstimator.ENSEMBLE == 'average':
+                # average
+                resized_heatMat[max(0, y):y + h, max(0, x):x + w, :] += heatMat[max(0, -y):, max(0, -x):, :]
+                resized_pafMat[max(0,y):y+h, max(0, x):x+w, :] += pafMat[max(0, -y):, max(0, -x):, :]
+                resized_cntMat[max(0,y):y+h, max(0, x):x+w, :] += 1
+            else:
+                # add up
+                resized_heatMat[max(0, y):y + h, max(0, x):x + w, :] = np.maximum(resized_heatMat[max(0, y):y + h, max(0, x):x + w, :], heatMat[max(0, -y):, max(0, -x):, :])
+                resized_pafMat[max(0,y):y+h, max(0, x):x+w, :] += pafMat[max(0, -y):, max(0, -x):, :]
+                resized_cntMat[max(0, y):y + h, max(0, x):x + w, :] += 1
+
+        if TfPoseEstimator.ENSEMBLE == 'average':
+            self.heatMat = resized_heatMat / resized_cntMat
+            self.pafMat = resized_pafMat / resized_cntMat
+        else:
+            self.heatMat = resized_heatMat
+            self.pafMat = resized_pafMat / (np.log(resized_cntMat) + 1)
+
+        humans = PoseEstimator.estimate(self.heatMat, self.pafMat)
+        #print('humans', humans)
+        return humans

posegait/src/network_base.py

@@ -0,0 +1,360 @@
+import sys
+
+import numpy as np
+import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
+import common
+
+DEFAULT_PADDING = 'SAME'
+
+
+_init_xavier = tf.contrib.layers.xavier_initializer()
+_init_norm = tf.truncated_normal_initializer(stddev=0.01)
+_init_zero = slim.init_ops.zeros_initializer()
+_l2_regularizer_00004 = tf.contrib.layers.l2_regularizer(0.00004)
+_l2_regularizer_convb = tf.contrib.layers.l2_regularizer(common.regularizer_conv)
+
+
+def layer(op):
+    '''
+    Decorator for composable network layers.
+    '''
+
+    def layer_decorated(self, *args, **kwargs):
+        # Automatically set a name if not provided.
+        name = kwargs.setdefault('name', self.get_unique_name(op.__name__))
+        # Figure out the layer inputs.
+        if len(self.terminals) == 0:
+            raise RuntimeError('No input variables found for layer %s.' % name)
+        elif len(self.terminals) == 1:
+            layer_input = self.terminals[0]
+        else:
+            layer_input = list(self.terminals)
+        # Perform the operation and get the output.
+        layer_output = op(self, layer_input, *args, **kwargs)
+        # Add to layer LUT.
+        self.layers[name] = layer_output
+        # This output is now the input for the next layer.
+        self.feed(layer_output)
+        # Return self for chained calls.
+        return self
+
+    return layer_decorated
+
+
+class BaseNetwork(object):
+    def __init__(self, inputs, trainable=True):
+        # The input nodes for this network
+        self.inputs = inputs
+        # The current list of terminal nodes
+        self.terminals = []
+        # Mapping from layer names to layers
+        self.layers = dict(inputs)
+        # If true, the resulting variables are set as trainable
+        self.trainable = trainable
+        # Switch variable for dropout
+        self.use_dropout = tf.placeholder_with_default(tf.constant(1.0),
+                                                       shape=[],
+                                                       name='use_dropout')
+        self.setup()
+
+    def setup(self):
+        '''Construct the network. '''
+        raise NotImplementedError('Must be implemented by the subclass.')
+
+    def load(self, data_path, session, ignore_missing=False):
+        '''
+        Load network weights.
+        data_path: The path to the numpy-serialized network weights
+        session: The current TensorFlow session
+        ignore_missing: If true, serialized weights for missing layers are ignored.
+        '''
+        data_dict = np.load(data_path, encoding='bytes').item()
+        for op_name in data_dict:
+            if isinstance(data_dict[op_name], np.ndarray):
+                if 'RMSProp' in op_name:
+                    continue
+                with tf.variable_scope('', reuse=True):
+                    var = tf.get_variable(op_name.replace(':0', ''))
+                    try:
+                        session.run(var.assign(data_dict[op_name]))
+                    except Exception as e:
+                        print(op_name)
+                        print(e)
+                        sys.exit(-1)
+            else:
+                with tf.variable_scope(op_name, reuse=True):
+                    for param_name, data in data_dict[op_name].items():
+                        try:
+                            var = tf.get_variable(param_name.decode("utf-8"))
+                            session.run(var.assign(data))
+                        except ValueError as e:
+                            print(e)
+                            if not ignore_missing:
+                                raise
+
+    def feed(self, *args):
+        '''Set the input(s) for the next operation by replacing the terminal nodes.
+        The arguments can be either layer names or the actual layers.
+        '''
+        assert len(args) != 0
+        self.terminals = []
+        for fed_layer in args:
+            try:
+                is_str = isinstance(fed_layer, basestring)
+            except NameError:
+                is_str = isinstance(fed_layer, str)
+            if is_str:
+                try:
+                    fed_layer = self.layers[fed_layer]
+                except KeyError:
+                    raise KeyError('Unknown layer name fed: %s' % fed_layer)
+            self.terminals.append(fed_layer)
+        return self
+
+    def get_output(self, name=None):
+        '''Returns the current network output.'''
+        if not name:
+            return self.terminals[-1]
+        else:
+            return self.layers[name]
+
+    def get_tensor(self, name):
+        return self.get_output(name)
+
+    def get_unique_name(self, prefix):
+        '''Returns an index-suffixed unique name for the given prefix.
+        This is used for auto-generating layer names based on the type-prefix.
+        '''
+        ident = sum(t.startswith(prefix) for t, _ in self.layers.items()) + 1
+        return '%s_%d' % (prefix, ident)
+
+    def make_var(self, name, shape, trainable=True):
+        '''Creates a new TensorFlow variable.'''
+        return tf.get_variable(name, shape, trainable=self.trainable & trainable, initializer=tf.contrib.layers.xavier_initializer())
+
+    def validate_padding(self, padding):
+        '''Verifies that the padding is one of the supported ones.'''
+        assert padding in ('SAME', 'VALID')
+
+    @layer
+    def normalize_vgg(self, input, name):
+        # normalize input -1.0 ~ 1.0
+        input = tf.divide(input, 255.0, name=name + '_divide')
+        input = tf.subtract(input, 0.5, name=name + '_subtract')
+        input = tf.multiply(input, 2.0, name=name + '_multiply')
+        return input
+
+    @layer
+    def normalize_mobilenet(self, input, name):
+        input = tf.divide(input, 255.0, name=name + '_divide')
+        input = tf.subtract(input, 0.5, name=name + '_subtract')
+        input = tf.multiply(input, 2.0, name=name + '_multiply')
+        return input
+
+    @layer
+    def normalize_nasnet(self, input, name):
+        input = tf.divide(input, 255.0, name=name + '_divide')
+        input = tf.subtract(input, 0.5, name=name + '_subtract')
+        input = tf.multiply(input, 2.0, name=name + '_multiply')
+        return input
+
+    @layer
+    def upsample(self, input, factor, name):
+        return tf.image.resize_bilinear(input, [int(input.get_shape()[1]) * factor, int(input.get_shape()[2]) * factor], name=name)
+
+    @layer
+    def separable_conv(self, input, k_h, k_w, c_o, stride, name, relu=True, set_bias=True):
+        with slim.arg_scope([slim.batch_norm], decay=0.999, fused=common.batchnorm_fused, is_training=self.trainable):
+            output = slim.separable_convolution2d(input,
+                                                  num_outputs=None,
+                                                  stride=stride,
+                                                  trainable=self.trainable,
+                                                  depth_multiplier=1.0,
+                                                  kernel_size=[k_h, k_w],
+                                                  # activation_fn=common.activation_fn if relu else None,
+                                                  activation_fn=None,
+                                                  # normalizer_fn=slim.batch_norm,
+                                                  weights_initializer=_init_xavier,
+                                                  # weights_initializer=_init_norm,
+                                                  weights_regularizer=_l2_regularizer_00004,
+                                                  biases_initializer=None,
+                                                  padding=DEFAULT_PADDING,
+                                                  scope=name + '_depthwise')
+
+            output = slim.convolution2d(output,
+                                        c_o,
+                                        stride=1,
+                                        kernel_size=[1, 1],
+                                        activation_fn=common.activation_fn if relu else None,
+                                        weights_initializer=_init_xavier,
+                                        # weights_initializer=_init_norm,
+                                        biases_initializer=_init_zero if set_bias else None,
+                                        normalizer_fn=slim.batch_norm,
+                                        trainable=self.trainable,
+                                        weights_regularizer=None,
+                                        scope=name + '_pointwise')
+
+        return output
+
+    @layer
+    def convb(self, input, k_h, k_w, c_o, stride, name, relu=True, set_bias=True, set_tanh=False):
+        with slim.arg_scope([slim.batch_norm], decay=0.999, fused=common.batchnorm_fused, is_training=self.trainable):
+            output = slim.convolution2d(input, c_o, kernel_size=[k_h, k_w],
+                                        stride=stride,
+                                        normalizer_fn=slim.batch_norm,
+                                        weights_regularizer=_l2_regularizer_convb,
+                                        weights_initializer=_init_xavier,
+                                        # weights_initializer=tf.truncated_normal_initializer(stddev=0.01),
+                                        biases_initializer=_init_zero if set_bias else None,
+                                        trainable=self.trainable,
+                                        activation_fn=common.activation_fn if relu else None,
+                                        scope=name)
+            if set_tanh:
+                output = tf.nn.sigmoid(output, name=name + '_extra_acv')
+        return output
+
+    @layer
+    def conv(self,
+             input,
+             k_h,
+             k_w,
+             c_o,
+             s_h,
+             s_w,
+             name,
+             relu=True,
+             padding=DEFAULT_PADDING,
+             group=1,
+             trainable=True,
+             biased=True):
+        # Verify that the padding is acceptable
+        self.validate_padding(padding)
+        # Get the number of channels in the input
+        c_i = int(input.get_shape()[-1])
+        # Verify that the grouping parameter is valid
+        assert c_i % group == 0
+        assert c_o % group == 0
+        # Convolution for a given input and kernel
+        convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding)
+        with tf.variable_scope(name) as scope:
+            kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o], trainable=self.trainable & trainable)
+            if group == 1:
+                # This is the common-case. Convolve the input without any further complications.
+                output = convolve(input, kernel)
+            else:
+                # Split the input into groups and then convolve each of them independently
+                input_groups = tf.split(3, group, input)
+                kernel_groups = tf.split(3, group, kernel)
+                output_groups = [convolve(i, k) for i, k in zip(input_groups, kernel_groups)]
+                # Concatenate the groups
+                output = tf.concat(3, output_groups)
+            # Add the biases
+            if biased:
+                biases = self.make_var('biases', [c_o], trainable=self.trainable & trainable)
+                output = tf.nn.bias_add(output, biases)
+
+            if relu:
+                # ReLU non-linearity
+                output = tf.nn.relu(output, name=scope.name)
+            return output
+
+    @layer
+    def relu(self, input, name):
+        return tf.nn.relu(input, name=name)
+
+    @layer
+    def max_pool(self, input, k_h, k_w, s_h, s_w, name, padding=DEFAULT_PADDING):
+        self.validate_padding(padding)
+        return tf.nn.max_pool(input,
+                              ksize=[1, k_h, k_w, 1],
+                              strides=[1, s_h, s_w, 1],
+                              padding=padding,
+                              name=name)
+
+    @layer
+    def avg_pool(self, input, k_h, k_w, s_h, s_w, name, padding=DEFAULT_PADDING):
+        self.validate_padding(padding)
+        return tf.nn.avg_pool(input,
+                              ksize=[1, k_h, k_w, 1],
+                              strides=[1, s_h, s_w, 1],
+                              padding=padding,
+                              name=name)
+
+    @layer
+    def lrn(self, input, radius, alpha, beta, name, bias=1.0):
+        return tf.nn.local_response_normalization(input,
+                                                  depth_radius=radius,
+                                                  alpha=alpha,
+                                                  beta=beta,
+                                                  bias=bias,
+                                                  name=name)
+
+    @layer
+    def concat(self, inputs, axis, name):
+        return tf.concat(axis=axis, values=inputs, name=name)
+
+    @layer
+    def add(self, inputs, name):
+        return tf.add_n(inputs, name=name)
+
+    @layer
+    def fc(self, input, num_out, name, relu=True):
+        with tf.variable_scope(name) as scope:
+            input_shape = input.get_shape()
+            if input_shape.ndims == 4:
+                # The input is spatial. Vectorize it first.
+                dim = 1
+                for d in input_shape[1:].as_list():
+                    dim *= d
+                feed_in = tf.reshape(input, [-1, dim])
+            else:
+                feed_in, dim = (input, input_shape[-1].value)
+            weights = self.make_var('weights', shape=[dim, num_out])
+            biases = self.make_var('biases', [num_out])
+            op = tf.nn.relu_layer if relu else tf.nn.xw_plus_b
+            fc = op(feed_in, weights, biases, name=scope.name)
+            return fc
+
+    @layer
+    def softmax(self, input, name):
+        input_shape = map(lambda v: v.value, input.get_shape())
+        if len(input_shape) > 2:
+            # For certain models (like NiN), the singleton spatial dimensions
+            # need to be explicitly squeezed, since they're not broadcast-able
+            # in TensorFlow's NHWC ordering (unlike Caffe's NCHW).
+            if input_shape[1] == 1 and input_shape[2] == 1:
+                input = tf.squeeze(input, squeeze_dims=[1, 2])
+            else:
+                raise ValueError('Rank 2 tensor input expected for softmax!')
+        return tf.nn.softmax(input, name=name)
+
+    @layer
+    def batch_normalization(self, input, name, scale_offset=True, relu=False):
+        # NOTE: Currently, only inference is supported
+        with tf.variable_scope(name) as scope:
+            shape = [input.get_shape()[-1]]
+            if scale_offset:
+                scale = self.make_var('scale', shape=shape)
+                offset = self.make_var('offset', shape=shape)
+            else:
+                scale, offset = (None, None)
+            output = tf.nn.batch_normalization(
+                input,
+                mean=self.make_var('mean', shape=shape),
+                variance=self.make_var('variance', shape=shape),
+                offset=offset,
+                scale=scale,
+                # TODO: This is the default Caffe batch norm eps
+                # Get the actual eps from parameters
+                variance_epsilon=1e-5,
+                name=name)
+            if relu:
+                output = tf.nn.relu(output)
+            return output
+
+    @layer
+    def dropout(self, input, keep_prob, name):
+        keep = 1 - self.use_dropout + (self.use_dropout * keep_prob)
+        return tf.nn.dropout(input, keep, name=name)

posegait/src/network_cmu.py

@@ -0,0 +1,163 @@
+import network_base
+import tensorflow as tf
+
+
+class CmuNetwork(network_base.BaseNetwork):
+    def setup(self):
+        (self.feed('image')
+             .normalize_vgg(name='preprocess')
+             .conv(3, 3, 64, 1, 1, name='conv1_1')
+             .conv(3, 3, 64, 1, 1, name='conv1_2')
+             .max_pool(2, 2, 2, 2, name='pool1_stage1')
+             .conv(3, 3, 128, 1, 1, name='conv2_1')
+             .conv(3, 3, 128, 1, 1, name='conv2_2')
+             .max_pool(2, 2, 2, 2, name='pool2_stage1')
+             .conv(3, 3, 256, 1, 1, name='conv3_1')
+             .conv(3, 3, 256, 1, 1, name='conv3_2')
+             .conv(3, 3, 256, 1, 1, name='conv3_3')
+             .conv(3, 3, 256, 1, 1, name='conv3_4')
+             .max_pool(2, 2, 2, 2, name='pool3_stage1')
+             .conv(3, 3, 512, 1, 1, name='conv4_1')
+             .conv(3, 3, 512, 1, 1, name='conv4_2')
+             .conv(3, 3, 256, 1, 1, name='conv4_3_CPM')
+             .conv(3, 3, 128, 1, 1, name='conv4_4_CPM')          # *****
+
+             .conv(3, 3, 128, 1, 1, name='conv5_1_CPM_L1')
+             .conv(3, 3, 128, 1, 1, name='conv5_2_CPM_L1')
+             .conv(3, 3, 128, 1, 1, name='conv5_3_CPM_L1')
+             .conv(1, 1, 512, 1, 1, name='conv5_4_CPM_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='conv5_5_CPM_L1'))
+
+        (self.feed('conv4_4_CPM')
+             .conv(3, 3, 128, 1, 1, name='conv5_1_CPM_L2')
+             .conv(3, 3, 128, 1, 1, name='conv5_2_CPM_L2')
+             .conv(3, 3, 128, 1, 1, name='conv5_3_CPM_L2')
+             .conv(1, 1, 512, 1, 1, name='conv5_4_CPM_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='conv5_5_CPM_L2'))
+
+        (self.feed('conv5_5_CPM_L1',
+                   'conv5_5_CPM_L2',
+                   'conv4_4_CPM')
+             .concat(3, name='concat_stage2')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage2_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage2_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage2_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage2_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage2_L1')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage2_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage2_L1'))
+
+        (self.feed('concat_stage2')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage2_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage2_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage2_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage2_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage2_L2')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage2_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage2_L2'))
+
+        (self.feed('Mconv7_stage2_L1',
+                   'Mconv7_stage2_L2',
+                   'conv4_4_CPM')
+             .concat(3, name='concat_stage3')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage3_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage3_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage3_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage3_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage3_L1')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage3_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage3_L1'))
+
+        (self.feed('concat_stage3')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage3_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage3_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage3_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage3_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage3_L2')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage3_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage3_L2'))
+
+        (self.feed('Mconv7_stage3_L1',
+                   'Mconv7_stage3_L2',
+                   'conv4_4_CPM')
+             .concat(3, name='concat_stage4')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage4_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage4_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage4_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage4_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage4_L1')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage4_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage4_L1'))
+
+        (self.feed('concat_stage4')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage4_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage4_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage4_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage4_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage4_L2')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage4_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage4_L2'))
+
+        (self.feed('Mconv7_stage4_L1',
+                   'Mconv7_stage4_L2',
+                   'conv4_4_CPM')
+             .concat(3, name='concat_stage5')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage5_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage5_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage5_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage5_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage5_L1')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage5_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage5_L1'))
+
+        (self.feed('concat_stage5')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage5_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage5_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage5_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage5_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage5_L2')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage5_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage5_L2'))
+
+        (self.feed('Mconv7_stage5_L1',
+                   'Mconv7_stage5_L2',
+                   'conv4_4_CPM')
+             .concat(3, name='concat_stage6')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage6_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage6_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage6_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage6_L1')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage6_L1')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage6_L1')
+             .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage6_L1'))
+
+        (self.feed('concat_stage6')
+             .conv(7, 7, 128, 1, 1, name='Mconv1_stage6_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv2_stage6_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv3_stage6_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv4_stage6_L2')
+             .conv(7, 7, 128, 1, 1, name='Mconv5_stage6_L2')
+             .conv(1, 1, 128, 1, 1, name='Mconv6_stage6_L2')
+             .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage6_L2'))
+
+        with tf.variable_scope('Openpose'):
+            (self.feed('Mconv7_stage6_L2',
+                       'Mconv7_stage6_L1')
+                 .concat(3, name='concat_stage7'))
+
+    def loss_l1_l2(self):
+         l1s = []
+         l2s = []
+         for layer_name in self.layers.keys():
+              if 'Mconv7' in layer_name and '_L1' in layer_name:
+                   l1s.append(self.layers[layer_name])
+              if 'Mconv7' in layer_name and '_L2' in layer_name:
+                   l2s.append(self.layers[layer_name])
+
+         return l1s, l2s
+
+    def loss_last(self):
+         return self.get_output('Mconv7_stage6_L1'), self.get_output('Mconv7_stage6_L2')
+
+    def restorable_variables(self):
+         return None

posegait/src/network_dsconv.py

@@ -0,0 +1,147 @@
+import network_base
+
+
+class DSConvNetwork(network_base.BaseNetwork):
+    def __init__(self, inputs, trainable=True, conv_width=1.0):
+        self.conv_width = conv_width
+        network_base.BaseNetwork.__init__(self, inputs, trainable)
+
+    def setup(self):
+        (self.feed('image')
+         .conv(3, 3, 64, 1, 1, name='conv1_1', trainable=False)
+         # .conv(3, 3, 64, 1, 1, name='conv1_2', trainable=True)     # TODO
+         .separable_conv(3, 3, round(self.conv_width * 64), 2, name='conv1_2')
+         # .max_pool(2, 2, 2, 2, name='pool1_stage1')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv2_1')
+         .separable_conv(3, 3, round(self.conv_width * 128), 2, name='conv2_2')
+         # .max_pool(2, 2, 2, 2, name='pool2_stage1')
+         .separable_conv(3, 3, round(self.conv_width * 256), 1, name='conv3_1')
+         .separable_conv(3, 3, round(self.conv_width * 256), 1, name='conv3_2')
+         .separable_conv(3, 3, round(self.conv_width * 256), 1, name='conv3_3')
+         .separable_conv(3, 3, round(self.conv_width * 256), 2, name='conv3_4')
+         # .max_pool(2, 2, 2, 2, name='pool3_stage1')
+         .separable_conv(3, 3, round(self.conv_width * 512), 1, name='conv4_1')
+         .separable_conv(3, 3, round(self.conv_width * 512), 1, name='conv4_2')
+         .separable_conv(3, 3, round(self.conv_width * 256), 1, name='conv4_3_CPM')
+         .separable_conv(3, 3, 128, 1, name='conv4_4_CPM')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_1_CPM_L1')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_2_CPM_L1')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_3_CPM_L1')
+         .conv(1, 1, 512, 1, 1, name='conv5_4_CPM_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='conv5_5_CPM_L1'))
+
+        (self.feed('conv4_4_CPM')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_1_CPM_L2')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_2_CPM_L2')
+         .separable_conv(3, 3, round(self.conv_width * 128), 1, name='conv5_3_CPM_L2')
+         .conv(1, 1, 512, 1, 1, name='conv5_4_CPM_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='conv5_5_CPM_L2'))
+
+        (self.feed('conv5_5_CPM_L1',
+                   'conv5_5_CPM_L2',
+                   'conv4_4_CPM')
+         .concat(3, name='concat_stage2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage2_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage2_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage2_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage2_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage2_L1')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage2_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage2_L1'))
+
+        (self.feed('concat_stage2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage2_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage2_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage2_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage2_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage2_L2')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage2_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage2_L2'))
+
+        (self.feed('Mconv7_stage2_L1',
+                   'Mconv7_stage2_L2',
+                   'conv4_4_CPM')
+         .concat(3, name='concat_stage3')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage3_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage3_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage3_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage3_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage3_L1')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage3_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage3_L1'))
+
+        (self.feed('concat_stage3')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage3_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage3_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage3_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage3_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage3_L2')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage3_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage3_L2'))
+
+        (self.feed('Mconv7_stage3_L1',
+                   'Mconv7_stage3_L2',
+                   'conv4_4_CPM')
+         .concat(3, name='concat_stage4')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage4_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage4_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage4_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage4_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage4_L1')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage4_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage4_L1'))
+
+        (self.feed('concat_stage4')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage4_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage4_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage4_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage4_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage4_L2')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage4_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage4_L2'))
+
+        (self.feed('Mconv7_stage4_L1',
+                   'Mconv7_stage4_L2',
+                   'conv4_4_CPM')
+         .concat(3, name='concat_stage5')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage5_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage5_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage5_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage5_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage5_L1')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage5_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage5_L1'))
+
+        (self.feed('concat_stage5')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage5_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage5_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage5_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage5_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage5_L2')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage5_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage5_L2'))
+
+        (self.feed('Mconv7_stage5_L1',
+                   'Mconv7_stage5_L2',
+                   'conv4_4_CPM')
+         .concat(3, name='concat_stage6')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage6_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage6_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage6_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage6_L1')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage6_L1')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage6_L1')
+         .conv(1, 1, 38, 1, 1, relu=False, name='Mconv7_stage6_L1'))
+
+        (self.feed('concat_stage6')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv1_stage6_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv2_stage6_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv3_stage6_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv4_stage6_L2')
+         .separable_conv(7, 7, round(self.conv_width * 128), 1, name='Mconv5_stage6_L2')
+         .conv(1, 1, 128, 1, 1, name='Mconv6_stage6_L2')
+         .conv(1, 1, 19, 1, 1, relu=False, name='Mconv7_stage6_L2'))
+
+        (self.feed('Mconv7_stage6_L2',
+                   'Mconv7_stage6_L1')
+         .concat(3, name='concat_stage7'))

posegait/src/network_mobilenet.py

@@ -0,0 +1,107 @@
+import tensorflow as tf
+
+import network_base
+
+
+class MobilenetNetwork(network_base.BaseNetwork):
+    def __init__(self, inputs, trainable=True, conv_width=1.0, conv_width2=None):
+        self.conv_width = conv_width
+        self.conv_width2 = conv_width2 if conv_width2 else conv_width
+        self.num_refine = 4
+        network_base.BaseNetwork.__init__(self, inputs, trainable)
+
+    def setup(self):
+        min_depth = 8
+        depth = lambda d: max(int(d * self.conv_width), min_depth)
+        depth2 = lambda d: max(int(d * self.conv_width2), min_depth)
+
+        with tf.variable_scope(None, 'MobilenetV1'):
+            (self.feed('image')
+             .convb(3, 3, depth(32), 2, name='Conv2d_0')
+             .separable_conv(3, 3, depth(64), 1, name='Conv2d_1')
+             .separable_conv(3, 3, depth(128), 2, name='Conv2d_2')
+             .separable_conv(3, 3, depth(128), 1, name='Conv2d_3')
+             .separable_conv(3, 3, depth(256), 2, name='Conv2d_4')
+             .separable_conv(3, 3, depth(256), 1, name='Conv2d_5')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_6')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_7')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_8')
+             # .separable_conv(3, 3, depth(512), 1, name='Conv2d_9')
+             # .separable_conv(3, 3, depth(512), 1, name='Conv2d_10')
+             # .separable_conv(3, 3, depth(512), 1, name='Conv2d_11')
+             # .separable_conv(3, 3, depth(1024), 2, name='Conv2d_12')
+             # .separable_conv(3, 3, depth(1024), 1, name='Conv2d_13')
+             )
+
+        (self.feed('Conv2d_1').max_pool(2, 2, 2, 2, name='Conv2d_1_pool'))
+        (self.feed('Conv2d_7').upsample(2, name='Conv2d_7_upsample'))
+
+        (self.feed('Conv2d_1_pool', 'Conv2d_3', 'Conv2d_7_upsample')
+            .concat(3, name='feat_concat'))
+
+        feature_lv = 'feat_concat'
+        with tf.variable_scope(None, 'Openpose'):
+            prefix = 'MConv_Stage1'
+            (self.feed(feature_lv)
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_1')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_2')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_3')
+             .separable_conv(1, 1, depth2(512), 1, name=prefix + '_L1_4')
+             .separable_conv(1, 1, 38, 1, relu=False, name=prefix + '_L1_5'))
+
+            (self.feed(feature_lv)
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_1')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_2')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_3')
+             .separable_conv(1, 1, depth2(512), 1, name=prefix + '_L2_4')
+             .separable_conv(1, 1, 19, 1, relu=False, name=prefix + '_L2_5'))
+
+            for stage_id in range(self.num_refine):
+                prefix_prev = 'MConv_Stage%d' % (stage_id + 1)
+                prefix = 'MConv_Stage%d' % (stage_id + 2)
+                (self.feed(prefix_prev + '_L1_5',
+                           prefix_prev + '_L2_5',
+                           feature_lv)
+                 .concat(3, name=prefix + '_concat')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L1_1')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L1_2')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L1_3')
+                 .separable_conv(1, 1, depth2(128), 1, name=prefix + '_L1_4')
+                 .separable_conv(1, 1, 38, 1, relu=False, name=prefix + '_L1_5'))
+
+                (self.feed(prefix + '_concat')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L2_1')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L2_2')
+                 .separable_conv(7, 7, depth2(128), 1, name=prefix + '_L2_3')
+                 .separable_conv(1, 1, depth2(128), 1, name=prefix + '_L2_4')
+                 .separable_conv(1, 1, 19, 1, relu=False, name=prefix + '_L2_5'))
+
+            # final result
+            (self.feed('MConv_Stage%d_L2_5' % self.get_refine_num(),
+                       'MConv_Stage%d_L1_5' % self.get_refine_num())
+             .concat(3, name='concat_stage7'))
+
+    def loss_l1_l2(self):
+        l1s = []
+        l2s = []
+        for layer_name in sorted(self.layers.keys()):
+            if '_L1_5' in layer_name:
+                l1s.append(self.layers[layer_name])
+            if '_L2_5' in layer_name:
+                l2s.append(self.layers[layer_name])
+
+        return l1s, l2s
+
+    def loss_last(self):
+        return self.get_output('MConv_Stage%d_L1_5' % self.get_refine_num()), \
+               self.get_output('MConv_Stage%d_L2_5' % self.get_refine_num())
+
+    def restorable_variables(self):
+        vs = {v.op.name: v for v in tf.global_variables() if
+              'MobilenetV1/Conv2d' in v.op.name and
+              'RMSProp' not in v.op.name and 'Momentum' not in v.op.name and 'Ada' not in v.op.name
+              }
+        return vs
+
+    def get_refine_num(self):
+        return self.num_refine + 1

posegait/src/network_mobilenet_thin.py

@@ -0,0 +1,104 @@
+import tensorflow as tf
+
+import network_base
+
+
+class MobilenetNetworkThin(network_base.BaseNetwork):
+    def __init__(self, inputs, trainable=True, conv_width=1.0, conv_width2=None):
+        self.conv_width = conv_width
+        self.conv_width2 = conv_width2 if conv_width2 else conv_width
+        network_base.BaseNetwork.__init__(self, inputs, trainable)
+
+    def setup(self):
+        min_depth = 8
+        depth = lambda d: max(int(d * self.conv_width), min_depth)
+        depth2 = lambda d: max(int(d * self.conv_width2), min_depth)
+
+        with tf.variable_scope(None, 'MobilenetV1'):
+            (self.feed('image')
+             .convb(3, 3, depth(32), 2, name='Conv2d_0')
+             .separable_conv(3, 3, depth(64), 1, name='Conv2d_1')
+             .separable_conv(3, 3, depth(128), 2, name='Conv2d_2')
+             .separable_conv(3, 3, depth(128), 1, name='Conv2d_3')
+             .separable_conv(3, 3, depth(256), 2, name='Conv2d_4')
+             .separable_conv(3, 3, depth(256), 1, name='Conv2d_5')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_6')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_7')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_8')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_9')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_10')
+             .separable_conv(3, 3, depth(512), 1, name='Conv2d_11')
+             # .separable_conv(3, 3, depth(1024), 2, name='Conv2d_12')
+             # .separable_conv(3, 3, depth(1024), 1, name='Conv2d_13')
+             )
+
+        (self.feed('Conv2d_3').max_pool(2, 2, 2, 2, name='Conv2d_3_pool'))
+
+        (self.feed('Conv2d_3_pool', 'Conv2d_7', 'Conv2d_11')
+            .concat(3, name='feat_concat'))
+
+        feature_lv = 'feat_concat'
+        with tf.variable_scope(None, 'Openpose'):
+            prefix = 'MConv_Stage1'
+            (self.feed(feature_lv)
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_1')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_2')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_3')
+             .separable_conv(1, 1, depth2(512), 1, name=prefix + '_L1_4')
+             .separable_conv(1, 1, 38, 1, relu=False, name=prefix + '_L1_5'))
+
+            (self.feed(feature_lv)
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_1')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_2')
+             .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_3')
+             .separable_conv(1, 1, depth2(512), 1, name=prefix + '_L2_4')
+             .separable_conv(1, 1, 19, 1, relu=False, name=prefix + '_L2_5'))
+
+            for stage_id in range(5):
+                prefix_prev = 'MConv_Stage%d' % (stage_id + 1)
+                prefix = 'MConv_Stage%d' % (stage_id + 2)
+                (self.feed(prefix_prev + '_L1_5',
+                           prefix_prev + '_L2_5',
+                           feature_lv)
+                 .concat(3, name=prefix + '_concat')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_1')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_2')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L1_3')
+                 .separable_conv(1, 1, depth2(128), 1, name=prefix + '_L1_4')
+                 .separable_conv(1, 1, 38, 1, relu=False, name=prefix + '_L1_5'))
+
+                (self.feed(prefix + '_concat')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_1')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_2')
+                 .separable_conv(3, 3, depth2(128), 1, name=prefix + '_L2_3')
+                 .separable_conv(1, 1, depth2(128), 1, name=prefix + '_L2_4')
+                 .separable_conv(1, 1, 19, 1, relu=False, name=prefix + '_L2_5'))
+
+            # final result
+            (self.feed('MConv_Stage6_L2_5',
+                       'MConv_Stage6_L1_5')
+             .concat(3, name='concat_stage7'))
+
+    def loss_l1_l2(self):
+        l1s = []
+        l2s = []
+        for layer_name in sorted(self.layers.keys()):
+            if '_L1_5' in layer_name:
+                l1s.append(self.layers[layer_name])
+            if '_L2_5' in layer_name:
+                l2s.append(self.layers[layer_name])
+
+        return l1s, l2s
+
+    def loss_last(self):
+        return self.get_output('MConv_Stage6_L1_5'), self.get_output('MConv_Stage6_L2_5')
+
+    def restorable_variables(self):
+        vs = {v.op.name: v for v in tf.global_variables() if
+              'MobilenetV1/Conv2d' in v.op.name and
+              # 'global_step' not in v.op.name and
+              # 'beta1_power' not in v.op.name and 'beta2_power' not in v.op.name and
+              'RMSProp' not in v.op.name and 'Momentum' not in v.op.name and
+              'Ada' not in v.op.name and 'Adam' not in v.op.name
+              }
+        return vs

posegait/src/networks.py

@@ -0,0 +1,93 @@
+import os
+
+import tensorflow as tf
+from network_mobilenet import MobilenetNetwork
+from network_mobilenet_thin import MobilenetNetworkThin
+
+from network_cmu import CmuNetwork
+
+
+def _get_base_path():
+    if not os.environ.get('OPENPOSE_MODEL', ''):
+        return './models'
+    return os.environ.get('OPENPOSE_MODEL')
+
+
+def get_network(type, placeholder_input, sess_for_load=None, trainable=True):
+    if type == 'mobilenet':
+        net = MobilenetNetwork({'image': placeholder_input}, conv_width=0.75, conv_width2=1.00, trainable=trainable)
+        pretrain_path = 'pretrained/mobilenet_v1_0.75_224_2017_06_14/mobilenet_v1_0.75_224.ckpt'
+        last_layer = 'MConv_Stage6_L{aux}_5'
+    elif type == 'mobilenet_fast':
+        net = MobilenetNetwork({'image': placeholder_input}, conv_width=0.5, conv_width2=0.5, trainable=trainable)
+        pretrain_path = 'pretrained/mobilenet_v1_0.75_224_2017_06_14/mobilenet_v1_0.75_224.ckpt'
+        last_layer = 'MConv_Stage6_L{aux}_5'
+    elif type == 'mobilenet_accurate':
+        net = MobilenetNetwork({'image': placeholder_input}, conv_width=1.00, conv_width2=1.00, trainable=trainable)
+        pretrain_path = 'pretrained/mobilenet_v1_1.0_224_2017_06_14/mobilenet_v1_1.0_224.ckpt'
+        last_layer = 'MConv_Stage6_L{aux}_5'
+
+    elif type == 'mobilenet_thin':
+        net = MobilenetNetworkThin({'image': placeholder_input}, conv_width=0.75, conv_width2=0.50, trainable=trainable)
+        pretrain_path = 'pretrained/mobilenet_v1_0.75_224_2017_06_14/mobilenet_v1_1.0_224.ckpt'
+        last_layer = 'MConv_Stage6_L{aux}_5'
+
+    elif type == 'cmu':
+        net = CmuNetwork({'image': placeholder_input}, trainable=trainable)
+        pretrain_path = 'numpy/openpose_coco.npy'
+        last_layer = 'Mconv7_stage6_L{aux}'
+    elif type == 'vgg':
+        net = CmuNetwork({'image': placeholder_input}, trainable=trainable)
+        pretrain_path = 'numpy/openpose_vgg16.npy'
+        last_layer = 'Mconv7_stage6_L{aux}'
+    else:
+        raise Exception('Invalid Mode.')
+
+    pretrain_path_full = os.path.join(_get_base_path(), pretrain_path)
+    if sess_for_load is not None:
+        if type == 'cmu' or type == 'vgg':
+            if not os.path.isfile(pretrain_path_full):
+                raise Exception('Model file doesn\'t exist, path=%s' % pretrain_path_full)
+            net.load(os.path.join(_get_base_path(), pretrain_path), sess_for_load)
+        else:
+            s = '%dx%d' % (placeholder_input.shape[2], placeholder_input.shape[1])
+            ckpts = {
+                'mobilenet': 'trained/mobilenet_%s/model-246038' % s,
+                'mobilenet_thin': 'trained/mobilenet_thin_%s/model-449003' % s,
+                'mobilenet_fast': 'trained/mobilenet_fast_%s/model-189000' % s,
+                'mobilenet_accurate': 'trained/mobilenet_accurate/model-170000'
+            }
+            ckpt_path = os.path.join(_get_base_path(), ckpts[type])
+            loader = tf.train.Saver()
+            try:
+                loader.restore(sess_for_load, ckpt_path)
+            except Exception as e:
+                raise Exception('Fail to load model files. \npath=%s\nerr=%s' % (ckpt_path, str(e)))
+
+    return net, pretrain_path_full, last_layer
+
+
+def get_graph_path(model_name):
+    print('Inside graph_path')
+    dyn_graph_path = {
+        'cmu': './models/graph/cmu/graph_opt.pb',
+        'mobilenet_thin': './models/graph/mobilenet_thin/graph_opt.pb'
+    }
+    graph_path = dyn_graph_path[model_name]
+    for path in (graph_path, os.path.join(os.path.dirname(os.path.abspath(__file__)), graph_path), os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', graph_path)):
+        if not os.path.isfile(path):
+            continue
+        return path
+    raise Exception('Graph file doesn\'t exist, path=%s' % graph_path)
+
+
+def model_wh(resolution_str):
+    print('model_wh')
+    print('chk1', resolution_str)
+    #resolution_str = int(float(resolution_str))
+    width, height = map(int, resolution_str.split('x'))
+    print('width', width)
+    print('height', height)
+    if width % 16 != 0 or height % 16 != 0:
+        raise Exception('Width and height should be multiples of 16. w=%d, h=%d' % (width, height))
+    return int(width), int(height)

posegait/src/pose_augment.py

@@ -0,0 +1,265 @@
+import math
+import random
+
+import cv2
+import numpy as np
+from tensorpack.dataflow.imgaug.geometry import RotationAndCropValid
+
+from common import CocoPart
+
+_network_w = 368
+_network_h = 368
+_scale = 2
+
+
+def set_network_input_wh(w, h):
+    global _network_w, _network_h
+    _network_w, _network_h = w, h
+
+
+def set_network_scale(scale):
+    global _scale
+    _scale = scale
+
+
+def pose_random_scale(meta):
+    scalew = random.uniform(0.8, 1.2)
+    scaleh = random.uniform(0.8, 1.2)
+    neww = int(meta.width * scalew)
+    newh = int(meta.height * scaleh)
+    dst = cv2.resize(meta.img, (neww, newh), interpolation=cv2.INTER_AREA)
+
+    # adjust meta data
+    adjust_joint_list = []
+    for joint in meta.joint_list:
+        adjust_joint = []
+        for point in joint:
+            if point[0] < -100 or point[1] < -100:
+                adjust_joint.append((-1000, -1000))
+                continue
+            # if point[0] <= 0 or point[1] <= 0 or int(point[0] * scalew + 0.5) > neww or int(
+            #                         point[1] * scaleh + 0.5) > newh:
+            #     adjust_joint.append((-1, -1))
+            #     continue
+            adjust_joint.append((int(point[0] * scalew + 0.5), int(point[1] * scaleh + 0.5)))
+        adjust_joint_list.append(adjust_joint)
+
+    meta.joint_list = adjust_joint_list
+    meta.width, meta.height = neww, newh
+    meta.img = dst
+    return meta
+
+
+def pose_resize_shortestedge_fixed(meta):
+    ratio_w = _network_w / meta.width
+    ratio_h = _network_h / meta.height
+    ratio = max(ratio_w, ratio_h)
+    return pose_resize_shortestedge(meta, int(min(meta.width * ratio + 0.5, meta.height * ratio + 0.5)))
+
+
+def pose_resize_shortestedge_random(meta):
+    ratio_w = _network_w / meta.width
+    ratio_h = _network_h / meta.height
+    ratio = min(ratio_w, ratio_h)
+    target_size = int(min(meta.width * ratio + 0.5, meta.height * ratio + 0.5))
+    target_size = int(target_size * random.uniform(0.95, 1.6))
+    # target_size = int(min(_network_w, _network_h) * random.uniform(0.7, 1.5))
+    return pose_resize_shortestedge(meta, target_size)
+
+
+def pose_resize_shortestedge(meta, target_size):
+    global _network_w, _network_h
+    img = meta.img
+
+    # adjust image
+    scale = target_size / min(meta.height, meta.width)
+    if meta.height < meta.width:
+        newh, neww = target_size, int(scale * meta.width + 0.5)
+    else:
+        newh, neww = int(scale * meta.height + 0.5), target_size
+
+    dst = cv2.resize(img, (neww, newh), interpolation=cv2.INTER_AREA)
+
+    pw = ph = 0
+    if neww < _network_w or newh < _network_h:
+        pw = max(0, (_network_w - neww) // 2)
+        ph = max(0, (_network_h - newh) // 2)
+        mw = (_network_w - neww) % 2
+        mh = (_network_h - newh) % 2
+        color = random.randint(0, 255)
+        dst = cv2.copyMakeBorder(dst, ph, ph+mh, pw, pw+mw, cv2.BORDER_CONSTANT, value=(color, 0, 0))
+
+    # adjust meta data
+    adjust_joint_list = []
+    for joint in meta.joint_list:
+        adjust_joint = []
+        for point in joint:
+            if point[0] < -100 or point[1] < -100:
+                adjust_joint.append((-1000, -1000))
+                continue
+            # if point[0] <= 0 or point[1] <= 0 or int(point[0]*scale+0.5) > neww or int(point[1]*scale+0.5) > newh:
+            #     adjust_joint.append((-1, -1))
+            #     continue
+            adjust_joint.append((int(point[0]*scale+0.5) + pw, int(point[1]*scale+0.5) + ph))
+        adjust_joint_list.append(adjust_joint)
+
+    meta.joint_list = adjust_joint_list
+    meta.width, meta.height = neww + pw * 2, newh + ph * 2
+    meta.img = dst
+    return meta
+
+
+def pose_crop_center(meta):
+    global _network_w, _network_h
+    target_size = (_network_w, _network_h)
+    x = (meta.width - target_size[0]) // 2 if meta.width > target_size[0] else 0
+    y = (meta.height - target_size[1]) // 2 if meta.height > target_size[1] else 0
+
+    return pose_crop(meta, x, y, target_size[0], target_size[1])
+
+
+def pose_crop_random(meta):
+    global _network_w, _network_h
+    target_size = (_network_w, _network_h)
+
+    for _ in range(50):
+        x = random.randrange(0, meta.width - target_size[0]) if meta.width > target_size[0] else 0
+        y = random.randrange(0, meta.height - target_size[1]) if meta.height > target_size[1] else 0
+
+        # check whether any face is inside the box to generate a reasonably-balanced datasets
+        for joint in meta.joint_list:
+            if x <= joint[CocoPart.Nose.value][0] < x + target_size[0] and y <= joint[CocoPart.Nose.value][1] < y + target_size[1]:
+                break
+
+    return pose_crop(meta, x, y, target_size[0], target_size[1])
+
+
+def pose_crop(meta, x, y, w, h):
+    # adjust image
+    target_size = (w, h)
+
+    img = meta.img
+    resized = img[y:y+target_size[1], x:x+target_size[0], :]
+
+    # adjust meta data
+    adjust_joint_list = []
+    for joint in meta.joint_list:
+        adjust_joint = []
+        for point in joint:
+            if point[0] < -100 or point[1] < -100:
+                adjust_joint.append((-1000, -1000))
+                continue
+            # if point[0] <= 0 or point[1] <= 0:
+            #     adjust_joint.append((-1000, -1000))
+            #     continue
+            new_x, new_y = point[0] - x, point[1] - y
+            # if new_x <= 0 or new_y <= 0 or new_x > target_size[0] or new_y > target_size[1]:
+            #     adjust_joint.append((-1, -1))
+            #     continue
+            adjust_joint.append((new_x, new_y))
+        adjust_joint_list.append(adjust_joint)
+
+    meta.joint_list = adjust_joint_list
+    meta.width, meta.height = target_size
+    meta.img = resized
+    return meta
+
+
+def pose_flip(meta):
+    r = random.uniform(0, 1.0)
+    if r > 0.5:
+        return meta
+
+    img = meta.img
+    img = cv2.flip(img, 1)
+
+    # flip meta
+    flip_list = [CocoPart.Nose, CocoPart.Neck, CocoPart.LShoulder, CocoPart.LElbow, CocoPart.LWrist, CocoPart.RShoulder, CocoPart.RElbow, CocoPart.RWrist,
+                 CocoPart.LHip, CocoPart.LKnee, CocoPart.LAnkle, CocoPart.RHip, CocoPart.RKnee, CocoPart.RAnkle,
+                 CocoPart.LEye, CocoPart.REye, CocoPart.LEar, CocoPart.REar, CocoPart.Background]
+    adjust_joint_list = []
+    for joint in meta.joint_list:
+        adjust_joint = []
+        for cocopart in flip_list:
+            point = joint[cocopart.value]
+            if point[0] < -100 or point[1] < -100:
+                adjust_joint.append((-1000, -1000))
+                continue
+            # if point[0] <= 0 or point[1] <= 0:
+            #     adjust_joint.append((-1, -1))
+            #     continue
+            adjust_joint.append((meta.width - point[0], point[1]))
+        adjust_joint_list.append(adjust_joint)
+
+    meta.joint_list = adjust_joint_list
+
+    meta.img = img
+    return meta
+
+
+def pose_rotation(meta):
+    deg = random.uniform(-15.0, 15.0)
+    img = meta.img
+
+    center = (img.shape[1] * 0.5, img.shape[0] * 0.5)       # x, y
+    rot_m = cv2.getRotationMatrix2D((int(center[0]), int(center[1])), deg, 1)
+    ret = cv2.warpAffine(img, rot_m, img.shape[1::-1], flags=cv2.INTER_AREA, borderMode=cv2.BORDER_CONSTANT)
+    if img.ndim == 3 and ret.ndim == 2:
+        ret = ret[:, :, np.newaxis]
+    neww, newh = RotationAndCropValid.largest_rotated_rect(ret.shape[1], ret.shape[0], deg)
+    neww = min(neww, ret.shape[1])
+    newh = min(newh, ret.shape[0])
+    newx = int(center[0] - neww * 0.5)
+    newy = int(center[1] - newh * 0.5)
+    # print(ret.shape, deg, newx, newy, neww, newh)
+    img = ret[newy:newy + newh, newx:newx + neww]
+
+    # adjust meta data
+    adjust_joint_list = []
+    for joint in meta.joint_list:
+        adjust_joint = []
+        for point in joint:
+            if point[0] < -100 or point[1] < -100:
+                adjust_joint.append((-1000, -1000))
+                continue
+            # if point[0] <= 0 or point[1] <= 0:
+            #     adjust_joint.append((-1, -1))
+            #     continue
+            x, y = _rotate_coord((meta.width, meta.height), (newx, newy), point, deg)
+            adjust_joint.append((x, y))
+        adjust_joint_list.append(adjust_joint)
+
+    meta.joint_list = adjust_joint_list
+    meta.width, meta.height = neww, newh
+    meta.img = img
+
+    return meta
+
+
+def _rotate_coord(shape, newxy, point, angle):
+    angle = -1 * angle / 180.0 * math.pi
+
+    ox, oy = shape
+    px, py = point
+
+    ox /= 2
+    oy /= 2
+
+    qx = math.cos(angle) * (px - ox) - math.sin(angle) * (py - oy)
+    qy = math.sin(angle) * (px - ox) + math.cos(angle) * (py - oy)
+
+    new_x, new_y = newxy
+
+    qx += ox - new_x
+    qy += oy - new_y
+
+    return int(qx + 0.5), int(qy + 0.5)
+
+
+def pose_to_img(meta_l):
+    global _network_w, _network_h, _scale
+    return [
+        meta_l[0].img.astype(np.float16),
+        meta_l[0].get_heatmap(target_size=(_network_w // _scale, _network_h // _scale)),
+        meta_l[0].get_vectormap(target_size=(_network_w // _scale, _network_h // _scale))
+    ]

posegait/src/pose_datamaster.py

@@ -0,0 +1,36 @@
+import argparse
+import logging
+import time
+
+from tensorpack.dataflow.remote import RemoteDataZMQ
+
+from pose_dataset import CocoPose
+
+logging.basicConfig(level=logging.DEBUG, format='[lmdb_dataset] %(asctime)s %(levelname)s %(message)s')
+
+if __name__ == '__main__':
+    """
+    Speed Test for Getting Input batches from other nodes
+    """
+    parser = argparse.ArgumentParser(description='Worker for preparing input batches.')
+    parser.add_argument('--listen', type=str, default='tcp://0.0.0.0:1027')
+    parser.add_argument('--show', type=bool, default=False)
+    args = parser.parse_args()
+
+    df = RemoteDataZMQ(args.listen)
+
+    logging.info('tcp queue start')
+    df.reset_state()
+    t = time.time()
+    for i, dp in enumerate(df.get_data()):
+        if i == 100:
+            break
+        logging.info('Input batch %d received.' % i)
+        if i == 0:
+            for d in dp:
+                logging.info('%d dp shape={}'.format(d.shape))
+
+        if args.show:
+            CocoPose.display_image(dp[0][0], dp[1][0], dp[2][0])
+
+    logging.info('Speed Test Done for 100 Batches in %f seconds.' % (time.time() - t))

posegait/src/pose_dataset.py

@@ -0,0 +1,485 @@
+import logging
+import math
+import multiprocessing
+import struct
+import sys
+import threading
+
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+from contextlib import contextmanager
+
+import os
+import random
+import requests
+import cv2
+import numpy as np
+import time
+
+import tensorflow as tf
+
+from tensorpack.dataflow import MultiThreadMapData
+from tensorpack.dataflow.image import MapDataComponent
+from tensorpack.dataflow.common import BatchData, MapData
+from tensorpack.dataflow.prefetch import PrefetchData
+from tensorpack.dataflow.base import RNGDataFlow, DataFlowTerminated
+
+from pycocotools.coco import COCO
+from pose_augment import pose_flip, pose_rotation, pose_to_img, pose_crop_random, \
+    pose_resize_shortestedge_random, pose_resize_shortestedge_fixed, pose_crop_center, pose_random_scale
+
+logging.getLogger("requests").setLevel(logging.WARNING)
+logger = logging.getLogger('pose_dataset')
+logger.setLevel(logging.INFO)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+mplset = False
+
+
+class CocoMetadata:
+    # __coco_parts = 57
+    __coco_parts = 19
+    __coco_vecs = list(zip(
+        [2, 9,  10, 2,  12, 13, 2, 3, 4, 3,  2, 6, 7, 6,  2, 1,  1,  15, 16],
+        [9, 10, 11, 12, 13, 14, 3, 4, 5, 17, 6, 7, 8, 18, 1, 15, 16, 17, 18]
+    ))
+
+    @staticmethod
+    def parse_float(four_np):
+        assert len(four_np) == 4
+        return struct.unpack('<f', bytes(four_np))[0]
+
+    @staticmethod
+    def parse_floats(four_nps, adjust=0):
+        assert len(four_nps) % 4 == 0
+        return [(CocoMetadata.parse_float(four_nps[x*4:x*4+4]) + adjust) for x in range(len(four_nps) // 4)]
+
+    def __init__(self, idx, img_url, img_meta, annotations, sigma):
+        self.idx = idx
+        self.img_url = img_url
+        self.img = None
+        self.sigma = sigma
+
+        self.height = int(img_meta['height'])
+        self.width = int(img_meta['width'])
+
+        joint_list = []
+        for ann in annotations:
+            if ann.get('num_keypoints', 0) == 0:
+                continue
+
+            kp = np.array(ann['keypoints'])
+            xs = kp[0::3]
+            ys = kp[1::3]
+            vs = kp[2::3]
+
+            joint_list.append([(x, y) if v >= 1 else (-1000, -1000) for x, y, v in zip(xs, ys, vs)])
+
+        self.joint_list = []
+        transform = list(zip(
+            [1, 6, 7, 9, 11, 6, 8, 10, 13, 15, 17, 12, 14, 16, 3, 2, 5, 4],
+            [1, 7, 7, 9, 11, 6, 8, 10, 13, 15, 17, 12, 14, 16, 3, 2, 5, 4]
+        ))
+        for prev_joint in joint_list:
+            new_joint = []
+            for idx1, idx2 in transform:
+                j1 = prev_joint[idx1-1]
+                j2 = prev_joint[idx2-1]
+
+                if j1[0] <= 0 or j1[1] <= 0 or j2[0] <= 0 or j2[1] <= 0:
+                    new_joint.append((-1000, -1000))
+                else:
+                    new_joint.append(((j1[0] + j2[0]) / 2, (j1[1] + j2[1]) / 2))
+
+            new_joint.append((-1000, -1000))
+            self.joint_list.append(new_joint)
+
+        # logger.debug('joint size=%d' % len(self.joint_list))
+
+    def get_heatmap(self, target_size):
+        heatmap = np.zeros((CocoMetadata.__coco_parts, self.height, self.width), dtype=np.float32)
+
+        for joints in self.joint_list:
+            for idx, point in enumerate(joints):
+                if point[0] < 0 or point[1] < 0:
+                    continue
+                CocoMetadata.put_heatmap(heatmap, idx, point, self.sigma)
+
+        heatmap = heatmap.transpose((1, 2, 0))
+
+        # background
+        heatmap[:, :, -1] = np.clip(1 - np.amax(heatmap, axis=2), 0.0, 1.0)
+
+        if target_size:
+            heatmap = cv2.resize(heatmap, target_size, interpolation=cv2.INTER_AREA)
+
+        return heatmap.astype(np.float16)
+
+    @staticmethod
+    def put_heatmap(heatmap, plane_idx, center, sigma):
+        center_x, center_y = center
+        _, height, width = heatmap.shape[:3]
+
+        th = 4.6052
+        delta = math.sqrt(th * 2)
+
+        x0 = int(max(0, center_x - delta * sigma))
+        y0 = int(max(0, center_y - delta * sigma))
+
+        x1 = int(min(width, center_x + delta * sigma))
+        y1 = int(min(height, center_y + delta * sigma))
+
+        for y in range(y0, y1):
+            for x in range(x0, x1):
+                d = (x - center_x) ** 2 + (y - center_y) ** 2
+                exp = d / 2.0 / sigma / sigma
+                if exp > th:
+                    continue
+                heatmap[plane_idx][y][x] = max(heatmap[plane_idx][y][x], math.exp(-exp))
+                heatmap[plane_idx][y][x] = min(heatmap[plane_idx][y][x], 1.0)
+
+    def get_vectormap(self, target_size):
+        vectormap = np.zeros((CocoMetadata.__coco_parts*2, self.height, self.width), dtype=np.float32)
+        countmap = np.zeros((CocoMetadata.__coco_parts, self.height, self.width), dtype=np.int16)
+        for joints in self.joint_list:
+            for plane_idx, (j_idx1, j_idx2) in enumerate(CocoMetadata.__coco_vecs):
+                j_idx1 -= 1
+                j_idx2 -= 1
+
+                center_from = joints[j_idx1]
+                center_to = joints[j_idx2]
+
+                if center_from[0] < -100 or center_from[1] < -100 or center_to[0] < -100 or center_to[1] < -100:
+                    continue
+
+                CocoMetadata.put_vectormap(vectormap, countmap, plane_idx, center_from, center_to)
+
+        vectormap = vectormap.transpose((1, 2, 0))
+        nonzeros = np.nonzero(countmap)
+        for p, y, x in zip(nonzeros[0], nonzeros[1], nonzeros[2]):
+            if countmap[p][y][x] <= 0:
+                continue
+            vectormap[y][x][p*2+0] /= countmap[p][y][x]
+            vectormap[y][x][p*2+1] /= countmap[p][y][x]
+
+        if target_size:
+            vectormap = cv2.resize(vectormap, target_size, interpolation=cv2.INTER_AREA)
+
+        return vectormap.astype(np.float16)
+
+    @staticmethod
+    def put_vectormap(vectormap, countmap, plane_idx, center_from, center_to, threshold=8):
+        _, height, width = vectormap.shape[:3]
+
+        vec_x = center_to[0] - center_from[0]
+        vec_y = center_to[1] - center_from[1]
+
+        min_x = max(0, int(min(center_from[0], center_to[0]) - threshold))
+        min_y = max(0, int(min(center_from[1], center_to[1]) - threshold))
+
+        max_x = min(width, int(max(center_from[0], center_to[0]) + threshold))
+        max_y = min(height, int(max(center_from[1], center_to[1]) + threshold))
+
+        norm = math.sqrt(vec_x ** 2 + vec_y ** 2)
+        if norm == 0:
+            return
+
+        vec_x /= norm
+        vec_y /= norm
+
+        for y in range(min_y, max_y):
+            for x in range(min_x, max_x):
+                bec_x = x - center_from[0]
+                bec_y = y - center_from[1]
+                dist = abs(bec_x * vec_y - bec_y * vec_x)
+
+                if dist > threshold:
+                    continue
+
+                countmap[plane_idx][y][x] += 1
+
+                vectormap[plane_idx*2+0][y][x] = vec_x
+                vectormap[plane_idx*2+1][y][x] = vec_y
+
+
+class CocoPose(RNGDataFlow):
+    @staticmethod
+    def display_image(inp, heatmap, vectmap, as_numpy=False):
+        global mplset
+        # if as_numpy and not mplset:
+        #     import matplotlib as mpl
+        #     mpl.use('Agg')
+        mplset = True
+        import matplotlib.pyplot as plt
+
+        fig = plt.figure()
+        a = fig.add_subplot(2, 2, 1)
+        a.set_title('Image')
+        plt.imshow(CocoPose.get_bgimg(inp))
+
+        a = fig.add_subplot(2, 2, 2)
+        a.set_title('Heatmap')
+        plt.imshow(CocoPose.get_bgimg(inp, target_size=(heatmap.shape[1], heatmap.shape[0])), alpha=0.5)
+        tmp = np.amax(heatmap, axis=2)
+        plt.imshow(tmp, cmap=plt.cm.gray, alpha=0.5)
+        plt.colorbar()
+
+        tmp2 = vectmap.transpose((2, 0, 1))
+        tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
+        tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)
+
+        a = fig.add_subplot(2, 2, 3)
+        a.set_title('Vectormap-x')
+        plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
+        plt.imshow(tmp2_odd, cmap=plt.cm.gray, alpha=0.5)
+        plt.colorbar()
+
+        a = fig.add_subplot(2, 2, 4)
+        a.set_title('Vectormap-y')
+        plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
+        plt.imshow(tmp2_even, cmap=plt.cm.gray, alpha=0.5)
+        plt.colorbar()
+
+        if not as_numpy:
+            plt.show()
+        else:
+            fig.canvas.draw()
+            data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+            data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+            fig.clear()
+            plt.close()
+            return data
+
+    @staticmethod
+    def get_bgimg(inp, target_size=None):
+        inp = cv2.cvtColor(inp.astype(np.uint8), cv2.COLOR_BGR2RGB)
+        if target_size:
+            inp = cv2.resize(inp, target_size, interpolation=cv2.INTER_AREA)
+        return inp
+
+    def __init__(self, path, img_path=None, is_train=True, decode_img=True, only_idx=-1):
+        self.is_train = is_train
+        self.decode_img = decode_img
+        self.only_idx = only_idx
+
+        if is_train:
+            whole_path = os.path.join(path, 'person_keypoints_train2017.json')
+        else:
+            whole_path = os.path.join(path, 'person_keypoints_val2017.json')
+        self.img_path = (img_path if img_path is not None else '') + ('train2017/' if is_train else 'val2017/')
+        self.coco = COCO(whole_path)
+
+        logger.info('%s dataset %d' % (path, self.size()))
+
+    def size(self):
+        return len(self.coco.imgs)
+
+    def get_data(self):
+        idxs = np.arange(self.size())
+        if self.is_train:
+            self.rng.shuffle(idxs)
+        else:
+            pass
+
+        keys = list(self.coco.imgs.keys())
+        for idx in idxs:
+            img_meta = self.coco.imgs[keys[idx]]
+            img_idx = img_meta['id']
+            ann_idx = self.coco.getAnnIds(imgIds=img_idx)
+
+            if 'http://' in self.img_path:
+                img_url = self.img_path + img_meta['file_name']
+            else:
+                img_url = os.path.join(self.img_path, img_meta['file_name'])
+
+            anns = self.coco.loadAnns(ann_idx)
+            meta = CocoMetadata(idx, img_url, img_meta, anns, sigma=8.0)
+
+            total_keypoints = sum([ann.get('num_keypoints', 0) for ann in anns])
+            if total_keypoints == 0 and random.uniform(0, 1) > 0.2:
+                continue
+
+            yield [meta]
+
+
+class MPIIPose(RNGDataFlow):
+    def __init__(self):
+        pass
+
+    def size(self):
+        pass
+
+    def get_data(self):
+        pass
+
+
+def read_image_url(metas):
+    for meta in metas:
+        img_str = None
+        if 'http://' in meta.img_url:
+            # print(meta.img_url)
+            for _ in range(10):
+                try:
+                    resp = requests.get(meta.img_url)
+                    if resp.status_code // 100 != 2:
+                        logger.warning('request failed code=%d url=%s' % (resp.status_code, meta.img_url))
+                        time.sleep(1.0)
+                        continue
+                    img_str = resp.content
+                    break
+                except Exception as e:
+                    logger.warning('request failed url=%s, err=%s' % (meta.img_url, str(e)))
+        else:
+            img_str = open(meta.img_url, 'rb').read()
+
+        if not img_str:
+            logger.warning('image not read, path=%s' % meta.img_url)
+            raise Exception()
+
+        nparr = np.fromstring(img_str, np.uint8)
+        meta.img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    return metas
+
+
+def get_dataflow(path, is_train, img_path=None):
+    ds = CocoPose(path, img_path, is_train)       # read data from lmdb
+    if is_train:
+        ds = MapData(ds, read_image_url)
+        ds = MapDataComponent(ds, pose_random_scale)
+        ds = MapDataComponent(ds, pose_rotation)
+        ds = MapDataComponent(ds, pose_flip)
+        ds = MapDataComponent(ds, pose_resize_shortestedge_random)
+        ds = MapDataComponent(ds, pose_crop_random)
+        ds = MapData(ds, pose_to_img)
+        # augs = [
+        #     imgaug.RandomApplyAug(imgaug.RandomChooseAug([
+        #         imgaug.GaussianBlur(max_size=3)
+        #     ]), 0.7)
+        # ]
+        # ds = AugmentImageComponent(ds, augs)
+        ds = PrefetchData(ds, 1000, multiprocessing.cpu_count() * 4)
+    else:
+        ds = MultiThreadMapData(ds, nr_thread=16, map_func=read_image_url, buffer_size=1000)
+        ds = MapDataComponent(ds, pose_resize_shortestedge_fixed)
+        ds = MapDataComponent(ds, pose_crop_center)
+        ds = MapData(ds, pose_to_img)
+        ds = PrefetchData(ds, 100, multiprocessing.cpu_count() // 4)
+
+    return ds
+
+
+def get_dataflow_batch(path, is_train, batchsize, img_path=None):
+    logger.info('dataflow img_path=%s' % img_path)
+    ds = get_dataflow(path, is_train, img_path=img_path)
+    ds = BatchData(ds, batchsize)
+    if is_train:
+        ds = PrefetchData(ds, 10, 2)
+    else:
+        ds = PrefetchData(ds, 50, 2)
+
+    return ds
+
+
+class DataFlowToQueue(threading.Thread):
+    def __init__(self, ds, placeholders, queue_size=5):
+        super(DataFlowToQueue).__init__()
+        self.daemon = True
+
+        self.ds = ds
+        self.placeholders = placeholders
+        self.queue = tf.FIFOQueue(queue_size, [ph.dtype for ph in placeholders], shapes=[ph.get_shape() for ph in placeholders])
+        self.op = self.queue.enqueue(placeholders)
+        self.close_op = self.queue.close(cancel_pending_enqueues=True)
+
+        self._coord = None
+        self._sess = None
+
+        self.last_dp = None
+
+    @contextmanager
+    def default_sess(self):
+        if self._sess:
+            with self._sess.as_default():
+                yield
+        else:
+            logger.warning("DataFlowToQueue {} wasn't under a default session!".format(self.name))
+            yield
+
+    def size(self):
+        return self.queue.size()
+
+    def start(self):
+        self._sess = tf.get_default_session()
+        super(DataFlowToQueue).start()
+
+    def set_coordinator(self, coord):
+        self._coord = coord
+
+    def run(self):
+        with self.default_sess():
+            try:
+                while not self._coord.should_stop():
+                    try:
+                        self.ds.reset_state()
+                        while True:
+                            for dp in self.ds.get_data():
+                                feed = dict(zip(self.placeholders, dp))
+                                self.op.run(feed_dict=feed)
+                                self.last_dp = dp
+                    except (tf.errors.CancelledError, tf.errors.OutOfRangeError, DataFlowTerminated):
+                        logger.error('err type1, placeholders={}'.format(self.placeholders))
+                        sys.exit(-1)
+                    except Exception as e:
+                        logger.error('err type2, err={}, placeholders={}'.format(str(e), self.placeholders))
+                        if isinstance(e, RuntimeError) and 'closed Session' in str(e):
+                            pass
+                        else:
+                            logger.exception("Exception in {}:{}".format(self.name, str(e)))
+                        sys.exit(-1)
+            except Exception as e:
+                logger.exception("Exception in {}:{}".format(self.name, str(e)))
+            finally:
+                try:
+                    self.close_op.run()
+                except Exception:
+                    pass
+                logger.info("{} Exited.".format(self.name))
+
+    def dequeue(self):
+        return self.queue.dequeue()
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_VISIBLE_DEVICES'] = ''
+
+    from src.pose_augment import set_network_input_wh
+    # set_network_input_wh(368, 368)
+    set_network_input_wh(480, 320)
+
+    df = get_dataflow('/root/coco/annotations', True, img_path='http://gpu-twg.kakaocdn.net/braincloud/COCO/')
+    # df = get_dataflow('/root/coco/annotations', False, img_path='http://gpu-twg.kakaocdn.net/braincloud/COCO/')
+
+    # TestDataSpeed(df).start()
+    # sys.exit(0)
+
+    with tf.Session() as sess:
+        df.reset_state()
+        t1 = time.time()
+        for idx, dp in enumerate(df.get_data()):
+            if idx == 0:
+                for d in dp:
+                    logger.info('%d dp shape={}'.format(d.shape))
+            print(time.time() - t1)
+            t1 = time.time()
+            CocoPose.display_image(dp[0], dp[1].astype(np.float32), dp[2].astype(np.float32))
+            print(dp[1].shape, dp[2].shape)
+            pass
+
+    logger.info('done')

posegait/src/pose_dataworker.py

@@ -0,0 +1,29 @@
+import argparse
+
+from tensorpack.dataflow.remote import send_dataflow_zmq
+
+from pose_dataset import get_dataflow_batch
+from pose_augment import set_network_input_wh, set_network_scale
+
+if __name__ == '__main__':
+    """
+    OpenPose Data Preparation might be a bottleneck for training.
+    You can run multiple workers to generate input batches in multi-nodes to make training process faster.
+    """
+    parser = argparse.ArgumentParser(description='Worker for preparing input batches.')
+    parser.add_argument('--datapath', type=str, default='/coco/annotations/')
+    parser.add_argument('--imgpath', type=str, default='/coco/')
+    parser.add_argument('--batchsize', type=int, default=64)
+    parser.add_argument('--train', type=bool, default=True)
+    parser.add_argument('--master', type=str, default='tcp://csi-cluster-gpu20.dakao.io:1027')
+    parser.add_argument('--input-width', type=int, default=368)
+    parser.add_argument('--input-height', type=int, default=368)
+    parser.add_argument('--scale-factor', type=int, default=2)
+    args = parser.parse_args()
+
+    set_network_input_wh(args.input_width, args.input_height)
+    set_network_scale(args.scale_factor)
+
+    df = get_dataflow_batch(args.datapath, args.train, args.batchsize, args.imgpath)
+
+    send_dataflow_zmq(df, args.master, hwm=10)

posegait/src/pose_stats.py

@@ -0,0 +1,74 @@
+from pose_dataset import CocoPose
+from tensorpack import imgaug
+from tensorpack.dataflow.common import MapDataComponent, MapData
+from tensorpack.dataflow.image import AugmentImageComponent
+
+from pose_augment import *
+
+
+def get_idx_hands_up():
+    from src.pose_augment import set_network_input_wh
+    set_network_input_wh(368, 368)
+
+    show_sample = True
+    db = CocoPoseLMDB('/data/public/rw/coco-pose-estimation-lmdb/', is_train=True, decode_img=show_sample)
+    db.reset_state()
+    total_cnt = 0
+    handup_cnt = 0
+    for idx, metas in enumerate(db.get_data()):
+        meta = metas[0]
+        if len(meta.joint_list) <= 0:
+            continue
+        body = meta.joint_list[0]
+        if body[CocoPart.Neck.value][1] <= 0:
+            continue
+        if body[CocoPart.LWrist.value][1] <= 0:
+            continue
+        if body[CocoPart.RWrist.value][1] <= 0:
+            continue
+
+        if body[CocoPart.Neck.value][1] > body[CocoPart.LWrist.value][1] or body[CocoPart.Neck.value][1] > body[CocoPart.RWrist.value][1]:
+            print(meta.idx)
+            handup_cnt += 1
+
+            if show_sample:
+                l1, l2, l3 = pose_to_img(metas)
+                CocoPose.display_image(l1, l2, l3)
+
+        total_cnt += 1
+
+    print('%d / %d' % (handup_cnt, total_cnt))
+
+
+def sample_augmentations():
+    ds = CocoPose('/data/public/rw/coco-pose-estimation-lmdb/', is_train=False, only_idx=0)
+    ds = MapDataComponent(ds, pose_random_scale)
+    ds = MapDataComponent(ds, pose_rotation)
+    ds = MapDataComponent(ds, pose_flip)
+    ds = MapDataComponent(ds, pose_resize_shortestedge_random)
+    ds = MapDataComponent(ds, pose_crop_random)
+    ds = MapData(ds, pose_to_img)
+    augs = [
+        imgaug.RandomApplyAug(imgaug.RandomChooseAug([
+            imgaug.GaussianBlur(3),
+            imgaug.SaltPepperNoise(white_prob=0.01, black_prob=0.01),
+            imgaug.RandomOrderAug([
+                imgaug.BrightnessScale((0.8, 1.2), clip=False),
+                imgaug.Contrast((0.8, 1.2), clip=False),
+                # imgaug.Saturation(0.4, rgb=True),
+            ]),
+        ]), 0.7),
+    ]
+    ds = AugmentImageComponent(ds, augs)
+
+    ds.reset_state()
+    for l1, l2, l3 in ds.get_data():
+        CocoPose.display_image(l1, l2, l3)
+
+
+if __name__ == '__main__':
+    # codes for tests
+    # get_idx_hands_up()
+
+    # show augmentation samples
+    sample_augmentations()

posegait/src/pt2-webcam3D.py

@@ -0,0 +1,105 @@
+"""
+This serve as our base openGL class.
+"""
+
+import numpy as np
+import pyqtgraph.opengl as gl
+from pyqtgraph.Qt import QtCore, QtGui
+import sys
+
+
+class Terrain(object):
+    def __init__(self):
+        """
+        Initialize the graphics window and mesh surface
+        """
+
+        # setup the view window
+        self.app = QtGui.QApplication(sys.argv)
+        self.window = gl.GLViewWidget()
+        self.window.setWindowTitle('Terrain')
+        self.window.setGeometry(0, 110, 1920, 1080)
+        self.window.setCameraPosition(distance=30, elevation=12)
+        self.window.show()
+
+        # constants and arrays
+        self.nsteps = 1
+        self.ypoints = np.arange(-20, 20 + self.nsteps, self.nsteps)
+        self.xpoints = np.arange(-20, 20 + self.nsteps, self.nsteps)
+        self.nfaces = len(self.ypoints)
+
+        # create the veritices array
+        verts, faces, colors = self.mesh()
+
+        self.mesh1 = gl.GLMeshItem(
+            faces=faces,
+            vertexes=verts,
+            faceColors=colors,
+            drawEdges=True,
+            smooth=False,
+        )
+        self.mesh1.setGLOptions('additive')
+        self.window.addItem(self.mesh1)
+
+    def mesh(self, height=2.5):
+
+        faces = []
+        colors = []
+        verts = np.array([
+            [
+                x, y, height * np.random.rand(1)
+            ] for xid, x in enumerate(self.xpoints) for yid, y in enumerate(self.ypoints)
+        ], dtype=np.float32)
+
+        for yid in range(self.nfaces - 1):
+            yoff = yid * self.nfaces
+            for xid in range(self.nfaces - 1):
+                faces.append([
+                    xid + yoff,
+                    xid + yoff + self.nfaces,
+                    xid + yoff + self.nfaces + 1,
+                ])
+                faces.append([
+                    xid + yoff,
+                    xid + yoff + 1,
+                    xid + yoff + self.nfaces + 1,
+                ])
+                colors.append([
+                    xid / self.nfaces, 1 - xid / self.nfaces, yid / self.nfaces, 0.7
+                ])
+                colors.append([
+                    xid / self.nfaces, 1 - xid / self.nfaces, yid / self.nfaces, 0.8
+                ])
+
+        faces = np.array(faces, dtype=np.uint32)
+        colors = np.array(colors, dtype=np.float32)
+
+        return verts, faces, colors
+
+    def update(self):
+        """
+        update the mesh and shift the noise each time
+        """
+        verts, faces, colors = self.mesh()
+        self.mesh1.setMeshData(vertexes=verts, faces=faces, faceColors=colors)
+
+    def start(self):
+        """
+        get the graphics window open and setup
+        """
+        if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
+            QtGui.QApplication.instance().exec_()
+
+    def animation(self, frametime=10):
+        """
+        calls the update method to run in a loop
+        """
+        timer = QtCore.QTimer()
+        timer.timeout.connect(self.update)
+        timer.start(frametime)
+        self.start()
+
+
+if __name__ == '__main__':
+    t = Terrain()
+    t.animation()

posegait/src/run_video.py

@@ -0,0 +1,61 @@
+import argparse
+import logging
+import time
+
+import cv2
+import numpy as np
+
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+
+logger = logging.getLogger('TfPoseEstimator-Video')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+fps_time = 0
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='tf-pose-estimation Video')
+    parser.add_argument('--video', type=str, default='')
+    parser.add_argument('--zoom', type=float, default=1.0)
+    parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
+    parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin')
+    parser.add_argument('--show-process', type=bool, default=False,
+                        help='for debug purpose, if enabled, speed for inference is dropped.')
+    args = parser.parse_args()
+
+    logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
+    w, h = model_wh(args.resolution)
+    e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
+    #logger.debug('cam read+')
+    #cam = cv2.VideoCapture(args.camera)
+    cap = cv2.VideoCapture(args.video)
+    #ret_val, image = cap.read()
+    #logger.info('cam image=%dx%d' % (image.shape[1], image.shape[0]))
+    if (cap.isOpened()== False):
+        print("Error opening video stream or file")
+    while(cap.isOpened()):
+        ret_val, image = cap.read()
+
+
+        humans = e.inference(image)
+        image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
+
+        #logger.debug('show+')
+        cv2.putText(image,
+                    "FPS: %f" % (1.0 / (time.time() - fps_time)),
+                    (10, 10),  cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+                    (0, 255, 0), 2)
+        cv2.imshow('tf-pose-estimation result', image)
+        fps_time = time.time()
+        if cv2.waitKey(1) == 27:
+            break
+
+
+    cv2.destroyAllWindows()
+logger.debug('finished+')

posegait/src/run_webcam.py

@@ -0,0 +1,80 @@
+import argparse
+import logging
+import time
+
+import cv2
+import numpy as np
+
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+
+logger = logging.getLogger('TfPoseEstimator-WebCam')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+fps_time = 0
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='tf-pose-estimation realtime webcam')
+    parser.add_argument('--camera', type=int, default=0)
+    parser.add_argument('--zoom', type=float, default=1.0)
+    parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
+    parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin')
+    parser.add_argument('--show-process', type=bool, default=False,
+                        help='for debug purpose, if enabled, speed for inference is dropped.')
+    args = parser.parse_args()
+
+    logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
+    w, h = model_wh(args.resolution)
+    e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
+    logger.debug('cam read+')
+    cam = cv2.VideoCapture(args.camera)
+    ret_val, image = cam.read()
+    #print('ret_val', ret_val)..Done
+    print('img', image)
+    logger.info('cam image=%dx%d' % (image.shape[1], image.shape[0]))
+
+    while True:
+        ret_val, image = cam.read()
+        #print('ret_val', ret_val)...Done
+        print('img', image)
+
+        logger.debug('image preprocess+')
+        if args.zoom < 1.0:
+            canvas = np.zeros_like(image)
+            img_scaled = cv2.resize(image, None, fx=args.zoom, fy=args.zoom, interpolation=cv2.INTER_LINEAR)
+            dx = (canvas.shape[1] - img_scaled.shape[1]) // 2
+            dy = (canvas.shape[0] - img_scaled.shape[0]) // 2
+            canvas[dy:dy + img_scaled.shape[0], dx:dx + img_scaled.shape[1]] = img_scaled
+            image = canvas
+        elif args.zoom > 1.0:
+            img_scaled = cv2.resize(image, None, fx=args.zoom, fy=args.zoom, interpolation=cv2.INTER_LINEAR)
+            dx = (img_scaled.shape[1] - image.shape[1]) // 2
+            dy = (img_scaled.shape[0] - image.shape[0]) // 2
+            image = img_scaled[dy:image.shape[0], dx:image.shape[1]]
+
+        print('img1', image)
+        logger.debug('image process+')
+        humans = e.inference(image)        
+
+        logger.debug('postprocess+')
+        image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
+        print('img2', image)
+
+        logger.debug('show+')
+        cv2.putText(image,
+                    "FPS: %f" % (1.0 / (time.time() - fps_time)),
+                    (10, 10),  cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+                    (0, 255, 0), 2)
+        cv2.imshow('tf-pose-estimation result', image)
+        fps_time = time.time()
+        if cv2.waitKey(1) == 27:
+            break
+        logger.debug('finished+')
+
+    cv2.destroyAllWindows()

posegait/src/runworking.py

@@ -0,0 +1,126 @@
+import argparse
+import logging
+import time
+import os
+import ast
+
+import common
+import cv2
+import numpy as np
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+import matplotlib.pyplot as plt
+from lifting.prob_model import Prob3dPose
+from lifting.draw import plot_pose
+
+logger = logging.getLogger('TfPoseEstimator')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+
+if __name__ == '__main__':
+    #print(os.getcwd())
+    #os.chdir('..')
+    #print(os.getcwd())
+    #import sys
+    #sys.exit(0)
+    #os.chdir('..')
+    parser = argparse.ArgumentParser(description='tf-pose-estimation run')
+    parser.add_argument('--image', type=str, default='dhoni.jpg')
+    parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
+    parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin')
+    parser.add_argument('--scales', type=str, default='[None]', help='for multiple scales, eg. [1.0, (1.1, 0.05)]')
+    args = parser.parse_args()
+    scales = ast.literal_eval(args.scales)
+
+    w, h = model_wh(args.resolution)
+    e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
+
+    # estimate human poses from a single image !
+    image = common.read_imgfile(args.image, None, None)
+    # image = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
+    t = time.time()
+    humans = e.inference(image, scales=scales)
+    elapsed = time.time() - t
+
+    logger.info('inference image: %s in %.4f seconds.' % (args.image, elapsed))
+
+    image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
+    cv2.imshow('tf-pose-estimation result', image)
+    #cv2.waitKey()
+
+    fig1 = plt.figure(1)
+    a = fig1.add_subplot(2, 2, 1)
+    a.set_title('Result')
+    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+    bgimg = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_BGR2RGB)
+    bgimg = cv2.resize(bgimg, (e.heatMat.shape[1], e.heatMat.shape[0]), interpolation=cv2.INTER_AREA)
+
+    # show network output
+    a = fig1.add_subplot(2, 2, 2)
+    plt.imshow(bgimg, alpha=0.5)
+    tmp = np.amax(e.heatMat[:, :, :-1], axis=2)
+    plt.imshow(tmp, cmap=plt.cm.gray, alpha=0.5)
+    plt.colorbar()
+
+    tmp2 = e.pafMat.transpose((2, 0, 1))
+    tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
+    tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)
+
+    a = fig1.add_subplot(2, 2, 3)
+    a.set_title('Vectormap-x')
+    #plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
+    #plt.imshow(tmp2_odd, cmap=plt.cm.gray, alpha=0.5)
+    #plt.colorbar()
+
+    a = fig1.add_subplot(2, 2, 4)
+    a.set_title('Vectormap-y')
+    #plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
+    plt.imshow(tmp2_even, cmap=plt.cm.gray, alpha=0.5)
+    plt.colorbar()
+    #plt.show()
+
+    #import sys
+    #sys.exit(0)
+
+    logger.info('3d lifting initialization.')
+    poseLifting = Prob3dPose('./lifting/models/prob_model_params.mat')
+
+    image_h, image_w = image.shape[:2]
+    standard_w = 640
+    standard_h = 480
+
+    fig2 = plt.figure(2)   
+    pose_2d_mpiis = []
+    visibilities = []
+    for human in humans:
+        pose_2d_mpii, visibility = common.MPIIPart.from_coco(human)
+        pose_2d_mpiis.append([(int(x * standard_w + 0.5), int(y * standard_h + 0.5)) for x, y in pose_2d_mpii])
+        visibilities.append(visibility)
+
+    pose_2d_mpiis = np.array(pose_2d_mpiis)
+    visibilities = np.array(visibilities)
+    transformed_pose2d, weights = poseLifting.transform_joints(pose_2d_mpiis, visibilities)
+    pose_3d = poseLifting.compute_3d(transformed_pose2d, weights)
+    lis_3d = pose_3d.tolist()
+    print('list_3d', lis_3d)
+    with open('C:\\Users\\carti\\Desktop\\tf-pose\\src\\img_3d.txt', 'a') as f:
+        for item in lis_3d:
+            f.write("%s\n" % item)
+    #print(pose_3d)
+    
+    pose_3dqt = np.array(pose_3d[0]).transpose()
+    
+    for point in pose_3dqt:
+        print(point)
+
+    for i, single_3d in enumerate(pose_3d):
+        plot_pose(single_3d)
+        plt.show()
+
+    pass

posegait/src/testwebcam3D.py

@@ -0,0 +1,178 @@
+"""
+This serve as our base openGL class.
+"""
+
+import numpy as np
+import pyqtgraph.opengl as gl
+import pyqtgraph as pg
+from pyqtgraph.Qt import QtCore, QtGui
+import sys
+import cv2
+import time
+import os
+import csv
+
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+from lifting.prob_model import Prob3dPose
+import common
+
+
+
+class Terrain(object):
+    def __init__(self):
+        """
+        Initialize the graphics window and mesh surface
+        """
+
+        # setup the view window
+        self.app = QtGui.QApplication(sys.argv)
+        self.window = gl.GLViewWidget()
+        self.window.setWindowTitle('Terrain')
+        self.window.setGeometry(0, 110, 1920, 1080)
+        self.window.setCameraPosition(distance=30, elevation=12)
+        self.window.show()
+
+        gx = gl.GLGridItem()
+        gy = gl.GLGridItem()
+        gz = gl.GLGridItem()
+        gx.rotate(90, 0, 1, 0)
+        gy.rotate(90, 1, 0, 0)
+        gx.translate(-10, 0, 0)
+        gy.translate(0, -10, 0)
+        gz.translate(0, 0, -10)
+        self.window.addItem(gx)
+        self.window.addItem(gy)
+        self.window.addItem(gz)
+
+        model = 'mobilenet_thin'
+        #model = '432x368'
+        camera = 0
+        self.lines = {}
+        self.connection = [
+            [0, 1], [1, 2], [2, 3], [0, 4], [4, 5], [5, 6],
+            [0, 7], [7, 8], [8, 9], [9, 10], [8, 11], [11, 12],
+            [12, 13], [8, 14], [14, 15], [15, 16]
+        ]
+        
+        resolution = '432x368'
+        w, h = model_wh(resolution)
+        self.e = TfPoseEstimator(get_graph_path(model), target_size=(w, h))
+        self.cam = cv2.VideoCapture(camera)
+        ret_val, image = self.cam.read()        
+        #print('ret_val', ret_val)
+        self.poseLifting = Prob3dPose('./src/lifting/models/prob_model_params.mat')
+        keypoints = self.mesh(image)
+        print('keypoints', keypoints)
+
+        self.points = gl.GLScatterPlotItem(
+            pos=keypoints,
+            color=pg.glColor((0, 255, 0)),
+            size=15
+        )
+        #print(keypoints)
+        self.window.addItem(self.points)
+        
+        for n, pts in enumerate(self.connection):
+            self.lines[n] = gl.GLLinePlotItem(
+                pos=np.array([keypoints[p] for p in pts]),
+                color=pg.glColor((0, 0, 255)),
+                width=3,
+                antialias=True
+            )
+            self.window.addItem(self.lines[n])
+
+
+    def mesh(self, image):
+        image_h, image_w = image.shape[:2]
+        width = 640
+        height = 480
+        pose_2d_mpiis = []
+        visibilities = []
+
+        humans = self.e.inference(image, scales=[None])
+        #print('humans3', humans)
+        
+            
+        for human in humans:
+            pose_2d_mpii, visibility = common.MPIIPart.from_coco(human)
+            pose_2d_mpiis.append(
+                [(int(x * width + 0.5), int(y * height + 0.5)) for x, y in pose_2d_mpii]
+            )
+            visibilities.append(visibility)
+
+        pose_2d_mpiis = np.array(pose_2d_mpiis)
+        #print('pose_2d_mpiis', pose_2d_mpiis)
+        visibilities = np.array(visibilities)
+        transformed_pose2d, weights = self.poseLifting.transform_joints(pose_2d_mpiis, visibilities)
+        #print('transformed_pose2d', transformed_pose2d)
+        pose_3d = self.poseLifting.compute_3d(transformed_pose2d, weights)
+        #print(type(pose_3d))
+        lis_3d = pose_3d.tolist()
+        print('list_3d', lis_3d)
+        #fil_pose = open('C:\\Users\\carti\\Desktop\\tf-pose\\src\\pose_3d.txt', 'w')
+        #fil_pose.write(lis_3d)
+        #fil_pose.close()
+        with open('C:\\Users\\carti\\Desktop\\tf-pose\\src\\pose_3d.txt', 'a') as f:
+            for item in lis_3d:
+                f.write("%s\n" % item)
+        #np.savetxt('C:\\Users\\carti\\Desktop\\tf-pose\\src\\pose_3da.txt', pose_3d)
+        
+        #with open('C:\\Users\\carti\\Desktop\\tf-pose\\src\\pose_3da.csv', 'a') as csvfile:
+        #    writer = csv.writer(csvfile, delimiter=",")
+        #    writer.writerow(lis_3d)
+        
+        #print(lis_3d[0])
+        #print('pose_3d', pose_3d.transpose())
+        #pose_chk1 = pose_3d.transpose()
+        #pose_chk = pose_chk1 / 80
+        #print('pose_3d', pose_chk)
+        keypoints = pose_3d[0].transpose()
+
+        return keypoints / 80
+
+    def update(self):
+        """
+        update the mesh and shift the noise each time
+        """
+        ret_val, image = self.cam.read()
+        try:
+            keypoints = self.mesh(image)
+        except AssertionError:
+            print('body not in image')
+        else:
+            self.points.setData(pos=keypoints)
+            
+            for n, pts in enumerate(self.connection):
+                self.lines[n].setData(
+                    pos=np.array([keypoints[p] for p in pts])
+                )
+
+    def start(self):
+        """
+        get the graphics window open and setup
+        """
+        if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
+            QtGui.QApplication.instance().exec_()
+            
+
+    def animation(self, frametime=10):
+        """
+        calls the update method to run in a loop
+        """
+        timer = QtCore.QTimer()
+        timer.timeout.connect(self.update)
+        timer.start(frametime)
+        self.start()
+        #Timer(5, self.exitfunc).start()        
+        #t = threading.Thread(target=self.listen)
+        #t.daemon = True
+        #t.start()
+        #time.sleep(3)
+        
+
+if __name__ == '__main__':
+    os.chdir('..')    
+    t = Terrain()
+    t.animation()        
+    

posegait/src/testwebcam3Dedited.py

@@ -0,0 +1,189 @@
+"""
+This serve as our base openGL class.
+"""
+
+import numpy as np
+import pyqtgraph.opengl as gl
+import pyqtgraph as pg
+from pyqtgraph.Qt import QtCore, QtGui
+import argparse
+import sys
+import logging
+import time
+import cv2
+import os
+
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+from lifting.prob_model import Prob3dPose
+import common
+
+logger = logging.getLogger('TfPoseEstimator-WebCam')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+class Terrain(object):
+    def __init__(self):
+
+        #print('path of pose est', os.getcwd())
+        parser = argparse.ArgumentParser(description='tf-pose-estimation realtime webcam')
+        parser.add_argument('--camera', type=int, default=0)
+        parser.add_argument('--zoom', type=float, default=1.0)
+        parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
+        parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin')
+        parser.add_argument('--show-process', type=bool, default=False,
+                            help='for debug purpose, if enabled, speed for inference is dropped.')
+        args = parser.parse_args()
+
+        logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
+        w, h = model_wh(args.resolution)
+        e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
+        logger.debug('cam read+')
+        cam = cv2.VideoCapture(args.camera)
+        ret_val, image = cam.read()
+        #print('ret_val', ret_val)..Done
+        print('img', image)
+        
+        while True:
+            ret_val, image = cam.read()
+            #print('ret_val', ret_val)...Done
+            print('img', image)
+
+            logger.debug('image preprocess+')
+            if args.zoom < 1.0:
+                canvas = np.zeros_like(image)
+                img_scaled = cv2.resize(image, None, fx=args.zoom, fy=args.zoom, interpolation=cv2.INTER_LINEAR)
+                dx = (canvas.shape[1] - img_scaled.shape[1]) // 2
+                dy = (canvas.shape[0] - img_scaled.shape[0]) // 2
+                canvas[dy:dy + img_scaled.shape[0], dx:dx + img_scaled.shape[1]] = img_scaled
+                image = canvas
+            elif args.zoom > 1.0:
+                img_scaled = cv2.resize(image, None, fx=args.zoom, fy=args.zoom, interpolation=cv2.INTER_LINEAR)
+                dx = (img_scaled.shape[1] - image.shape[1]) // 2
+                dy = (img_scaled.shape[0] - image.shape[0]) // 2
+                image = img_scaled[dy:image.shape[0], dx:image.shape[1]]
+
+            print('img1', image)
+            logger.debug('image process+')
+            humans = e.inference(image)        
+
+            logger.debug('postprocess+')
+            image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
+            print('img2', image)
+
+            logger.debug('show+')
+            
+            fps_time = 0
+            
+            cv2.putText(image,
+                        "FPS: %f" % (1.0 / (time.time() - fps_time)),
+                        (10, 10),  cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+                        (0, 255, 0), 2)
+            cv2.imshow('tf-pose-estimation result', image)
+            fps_time = time.time()
+            if cv2.waitKey(1) == 27:
+                break
+            logger.debug('finished+')
+
+            #cv2.destroyAllWindows()    
+
+            # setup the view window
+            app = QtGui.QApplication(sys.argv)
+            window = gl.GLViewWidget()
+            window.setWindowTitle('Terrain')
+            window.setGeometry(0, 110, 1920, 1080)
+            window.setCameraPosition(distance=30, elevation=12)
+            window.show()
+
+            gx = gl.GLGridItem()
+            gy = gl.GLGridItem()
+            gz = gl.GLGridItem()
+            gx.rotate(90, 0, 1, 0)
+            gy.rotate(90, 1, 0, 0)
+            gx.translate(-10, 0, 0)
+            gy.translate(0, -10, 0)
+            gz.translate(0, 0, -10)
+            window.addItem(gx)
+            window.addItem(gy)
+            window.addItem(gz)            
+            
+            poseLifting = Prob3dPose('./lifting/models/prob_model_params.mat')
+            
+            #keypoints = self.mesh(image)
+
+            #points = gl.GLScatterPlotItem(
+            #    pos=keypoints,
+            #    color=pg.glColor((0, 255, 0)),
+            #    size=15
+            #)
+            #window.addItem(points)    
+            
+    def mesh(self, image):
+        image_h, image_w = image.shape[:2]
+        width = 640
+        height = 480
+        pose_2d_mpiis = []
+        visibilities = []
+        model = 'mobilenet_thin'
+        self.e = TfPoseEstimator(get_graph_path(model), target_size=(width, height))
+        humans = self.e.inference(image, scales=[None])
+
+        for human in humans:
+            pose_2d_mpii, visibility = common.MPIIPart.from_coco(human)
+            pose_2d_mpiis.append(
+                [(int(x * width + 0.5), int(y * height + 0.5)) for x, y in pose_2d_mpii]
+            )
+            visibilities.append(visibility)
+
+        pose_2d_mpiis = np.array(pose_2d_mpiis)
+        visibilities = np.array(visibilities)
+        
+        poseLifting = Prob3dPose('./lifting/models/prob_model_params.mat')
+        
+        transformed_pose2d, weights = poseLifting.transform_joints(pose_2d_mpiis, visibilities)
+        pose_3d = self.poseLifting.compute_3d(transformed_pose2d, weights)
+
+        keypoints = pose_3d[0].transpose()
+
+        return keypoints / 80        
+
+    def update(self):
+        """
+        update the mesh and shift the noise each time
+        """
+        ret_val, image = self.cam.read()
+        try:
+            keypoints = self.mesh(image)
+        except AssertionError:
+            print('body not in image')
+        else:
+            self.points.setData(pos=keypoints)
+
+    def start(self):
+        """
+        get the graphics window open and setup
+        """
+        if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
+            QtGui.QApplication.instance().exec_()
+
+    def animation(self, frametime=10):
+        """
+        calls the update method to run in a loop
+        """
+        timer = QtCore.QTimer()
+        timer.timeout.connect(self.update)
+        timer.start(frametime)
+        self.start()
+    
+if __name__ == '__main__':
+    #os.chdir('..')
+    #print('path', os.getcwd())
+    a_long_time = 5
+    time.sleep(a_long_time)
+    TIMEOUT = 15
+    t = Terrain()
+    t.animation()

posegait/src/train.py

@@ -0,0 +1,282 @@
+import matplotlib as mpl
+mpl.use('Agg')      # training mode, no screen should be open. (It will block training loop)
+
+import argparse
+import logging
+import os
+import time
+
+
+import cv2
+import numpy as np
+import tensorflow as tf
+from tqdm import tqdm
+from tensorpack.dataflow.remote import RemoteDataZMQ
+
+from pose_dataset import get_dataflow_batch, DataFlowToQueue, CocoPose
+from pose_augment import set_network_input_wh, set_network_scale
+from common import get_sample_images
+from networks import get_network
+
+logger = logging.getLogger('train')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Training codes for Openpose using Tensorflow')
+    parser.add_argument('--model', default='mobilenet', help='model name')
+    parser.add_argument('--datapath', type=str, default='/root/coco/annotations')
+    parser.add_argument('--imgpath', type=str, default='/root/coco/')
+    parser.add_argument('--batchsize', type=int, default=96)
+    parser.add_argument('--gpus', type=int, default=1)
+    parser.add_argument('--max-epoch', type=int, default=30)
+    parser.add_argument('--lr', type=str, default='0.01')
+    parser.add_argument('--modelpath', type=str, default='/data/private/tf-openpose-models-2018-1/')
+    parser.add_argument('--logpath', type=str, default='/data/private/tf-openpose-log-2018-1/')
+    parser.add_argument('--checkpoint', type=str, default='')
+    parser.add_argument('--tag', type=str, default='')
+    parser.add_argument('--remote-data', type=str, default='', help='eg. tcp://0.0.0.0:1027')
+
+    parser.add_argument('--input-width', type=int, default=368)
+    parser.add_argument('--input-height', type=int, default=368)
+    args = parser.parse_args()
+
+    if args.gpus <= 0:
+        raise Exception('gpus <= 0')
+
+    # define input placeholder
+    set_network_input_wh(args.input_width, args.input_height)
+    scale = 4
+
+    if args.model in ['cmu', 'vgg', 'mobilenet_thin', 'mobilenet_try', 'mobilenet_try2', 'mobilenet_try3', 'hybridnet_try']:
+        scale = 8
+
+    set_network_scale(scale)
+    output_w, output_h = args.input_width // scale, args.input_height // scale
+
+    logger.info('define model+')
+    with tf.device(tf.DeviceSpec(device_type="GPU", device_index=0)):
+        input_node = tf.placeholder(tf.float32, shape=(args.batchsize, args.input_height, args.input_width, 3), name='image')
+        vectmap_node = tf.placeholder(tf.float32, shape=(args.batchsize, output_h, output_w, 38), name='vectmap')
+        heatmap_node = tf.placeholder(tf.float32, shape=(args.batchsize, output_h, output_w, 19), name='heatmap')
+
+        # prepare data
+        if not args.remote_data:
+            df = get_dataflow_batch(args.datapath, True, args.batchsize, img_path=args.imgpath)
+        else:
+            # transfer inputs from ZMQ
+            df = RemoteDataZMQ(args.remote_data, hwm=3)
+        enqueuer = DataFlowToQueue(df, [input_node, heatmap_node, vectmap_node], queue_size=100)
+        q_inp, q_heat, q_vect = enqueuer.dequeue()
+
+    df_valid = get_dataflow_batch(args.datapath, False, args.batchsize, img_path=args.imgpath)
+    df_valid.reset_state()
+    validation_cache = []
+
+    val_image = get_sample_images(args.input_width, args.input_height)
+    logger.info('tensorboard val image: %d' % len(val_image))
+    logger.info(q_inp)
+    logger.info(q_heat)
+    logger.info(q_vect)
+
+    # define model for multi-gpu
+    q_inp_split, q_heat_split, q_vect_split = tf.split(q_inp, args.gpus), tf.split(q_heat, args.gpus), tf.split(q_vect, args.gpus)
+
+    output_vectmap = []
+    output_heatmap = []
+    losses = []
+    last_losses_l1 = []
+    last_losses_l2 = []
+    outputs = []
+    for gpu_id in range(args.gpus):
+        with tf.device(tf.DeviceSpec(device_type="GPU", device_index=gpu_id)):
+            with tf.variable_scope(tf.get_variable_scope(), reuse=(gpu_id > 0)):
+                net, pretrain_path, last_layer = get_network(args.model, q_inp_split[gpu_id])
+                vect, heat = net.loss_last()
+                output_vectmap.append(vect)
+                output_heatmap.append(heat)
+                outputs.append(net.get_output())
+
+                l1s, l2s = net.loss_l1_l2()
+                for idx, (l1, l2) in enumerate(zip(l1s, l2s)):
+                    loss_l1 = tf.nn.l2_loss(tf.concat(l1, axis=0) - q_vect_split[gpu_id], name='loss_l1_stage%d_tower%d' % (idx, gpu_id))
+                    loss_l2 = tf.nn.l2_loss(tf.concat(l2, axis=0) - q_heat_split[gpu_id], name='loss_l2_stage%d_tower%d' % (idx, gpu_id))
+                    losses.append(tf.reduce_mean([loss_l1, loss_l2]))
+
+                last_losses_l1.append(loss_l1)
+                last_losses_l2.append(loss_l2)
+
+    outputs = tf.concat(outputs, axis=0)
+
+    with tf.device(tf.DeviceSpec(device_type="GPU", device_index=gpu_id)):
+        # define loss
+        total_loss = tf.reduce_sum(losses) / args.batchsize
+        total_loss_ll_paf = tf.reduce_sum(last_losses_l1) / args.batchsize
+        total_loss_ll_heat = tf.reduce_sum(last_losses_l2) / args.batchsize
+        total_loss_ll = tf.reduce_mean([total_loss_ll_paf, total_loss_ll_heat])
+
+        # define optimizer
+        step_per_epoch = 121745 // args.batchsize
+        global_step = tf.Variable(0, trainable=False)
+        if ',' not in args.lr:
+            starter_learning_rate = float(args.lr)
+            learning_rate = tf.train.exponential_decay(starter_learning_rate, global_step,
+                                                       decay_steps=10000, decay_rate=0.33, staircase=True)
+        else:
+            lrs = [float(x) for x in args.lr.split(',')]
+            boundaries = [step_per_epoch * 5 * i for i, _ in range(len(lrs)) if i > 0]
+            learning_rate = tf.train.piecewise_constant(global_step, boundaries, lrs)
+
+    # optimizer = tf.train.RMSPropOptimizer(learning_rate, decay=0.0005, momentum=0.9, epsilon=1e-10)
+    optimizer = tf.train.AdamOptimizer(learning_rate, epsilon=1e-8)
+    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+    with tf.control_dependencies(update_ops):
+        train_op = optimizer.minimize(total_loss, global_step, colocate_gradients_with_ops=True)
+    logger.info('define model-')
+
+    # define summary
+    tf.summary.scalar("loss", total_loss)
+    tf.summary.scalar("loss_lastlayer", total_loss_ll)
+    tf.summary.scalar("loss_lastlayer_paf", total_loss_ll_paf)
+    tf.summary.scalar("loss_lastlayer_heat", total_loss_ll_heat)
+    tf.summary.scalar("queue_size", enqueuer.size())
+    merged_summary_op = tf.summary.merge_all()
+
+    valid_loss = tf.placeholder(tf.float32, shape=[])
+    valid_loss_ll = tf.placeholder(tf.float32, shape=[])
+    valid_loss_ll_paf = tf.placeholder(tf.float32, shape=[])
+    valid_loss_ll_heat = tf.placeholder(tf.float32, shape=[])
+    sample_train = tf.placeholder(tf.float32, shape=(4, 640, 640, 3))
+    sample_valid = tf.placeholder(tf.float32, shape=(12, 640, 640, 3))
+    train_img = tf.summary.image('training sample', sample_train, 4)
+    valid_img = tf.summary.image('validation sample', sample_valid, 12)
+    valid_loss_t = tf.summary.scalar("loss_valid", valid_loss)
+    valid_loss_ll_t = tf.summary.scalar("loss_valid_lastlayer", valid_loss_ll)
+    merged_validate_op = tf.summary.merge([train_img, valid_img, valid_loss_t, valid_loss_ll_t])
+
+    saver = tf.train.Saver(max_to_keep=100)
+    config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
+    with tf.Session(config=config) as sess:
+        training_name = '{}_batch:{}_lr:{}_gpus:{}_{}x{}_{}'.format(
+            args.model,
+            args.batchsize,
+            args.lr,
+            args.gpus,
+            args.input_width, args.input_height,
+            args.tag
+        )
+        logger.info('model weights initialization')
+        sess.run(tf.global_variables_initializer())
+
+        if args.checkpoint:
+            logger.info('Restore from checkpoint...')
+            # loader = tf.train.Saver(net.restorable_variables())
+            # loader.restore(sess, tf.train.latest_checkpoint(args.checkpoint))
+            saver.restore(sess, tf.train.latest_checkpoint(args.checkpoint))
+            logger.info('Restore from checkpoint...Done')
+        elif pretrain_path:
+            logger.info('Restore pretrained weights...')
+            if '.ckpt' in pretrain_path:
+                loader = tf.train.Saver(net.restorable_variables())
+                loader.restore(sess, pretrain_path)
+            elif '.npy' in pretrain_path:
+                net.load(pretrain_path, sess, False)
+            logger.info('Restore pretrained weights...Done')
+
+        logger.info('prepare file writer')
+        file_writer = tf.summary.FileWriter(args.logpath + training_name, sess.graph)
+
+        logger.info('prepare coordinator')
+        coord = tf.train.Coordinator()
+        enqueuer.set_coordinator(coord)
+        enqueuer.start()
+
+        logger.info('Training Started.')
+        time_started = time.time()
+        last_gs_num = last_gs_num2 = 0
+        initial_gs_num = sess.run(global_step)
+
+        while True:
+            _, gs_num = sess.run([train_op, global_step])
+
+            if gs_num > step_per_epoch * args.max_epoch:
+                break
+
+            if gs_num - last_gs_num >= 100:
+                train_loss, train_loss_ll, train_loss_ll_paf, train_loss_ll_heat, lr_val, summary, queue_size = sess.run([total_loss, total_loss_ll, total_loss_ll_paf, total_loss_ll_heat, learning_rate, merged_summary_op, enqueuer.size()])
+
+                # log of training loss / accuracy
+                batch_per_sec = (gs_num - initial_gs_num) / (time.time() - time_started)
+                logger.info('epoch=%.2f step=%d, %0.4f examples/sec lr=%f, loss=%g, loss_ll=%g, loss_ll_paf=%g, loss_ll_heat=%g, q=%d' % (gs_num / step_per_epoch, gs_num, batch_per_sec * args.batchsize, lr_val, train_loss, train_loss_ll, train_loss_ll_paf, train_loss_ll_heat, queue_size))
+                last_gs_num = gs_num
+
+                file_writer.add_summary(summary, gs_num)
+
+            if gs_num - last_gs_num2 >= 1000:
+                # save weights
+                saver.save(sess, os.path.join(args.modelpath, training_name, 'model'), global_step=global_step)
+
+                average_loss = average_loss_ll = average_loss_ll_paf = average_loss_ll_heat = 0
+                total_cnt = 0
+
+                if len(validation_cache) == 0:
+                    for images_test, heatmaps, vectmaps in tqdm(df_valid.get_data()):
+                        validation_cache.append((images_test, heatmaps, vectmaps))
+                    df_valid.reset_state()
+                    del df_valid
+                    df_valid = None
+
+                # log of test accuracy
+                for images_test, heatmaps, vectmaps in validation_cache:
+                    lss, lss_ll, lss_ll_paf, lss_ll_heat, vectmap_sample, heatmap_sample = sess.run(
+                        [total_loss, total_loss_ll, total_loss_ll_paf, total_loss_ll_heat, output_vectmap, output_heatmap],
+                        feed_dict={q_inp: images_test, q_vect: vectmaps, q_heat: heatmaps}
+                    )
+                    average_loss += lss * len(images_test)
+                    average_loss_ll += lss_ll * len(images_test)
+                    average_loss_ll_paf += lss_ll_paf * len(images_test)
+                    average_loss_ll_heat += lss_ll_heat * len(images_test)
+                    total_cnt += len(images_test)
+
+                logger.info('validation(%d) %s loss=%f, loss_ll=%f, loss_ll_paf=%f, loss_ll_heat=%f' % (total_cnt, training_name, average_loss / total_cnt, average_loss_ll / total_cnt, average_loss_ll_paf / total_cnt, average_loss_ll_heat / total_cnt))
+                last_gs_num2 = gs_num
+
+                sample_image = [enqueuer.last_dp[0][i] for i in range(4)]
+                outputMat = sess.run(
+                    outputs,
+                    feed_dict={q_inp: np.array((sample_image + val_image)*(args.batchsize // 16))}
+                )
+                pafMat, heatMat = outputMat[:, :, :, 19:], outputMat[:, :, :, :19]
+
+                sample_results = []
+                for i in range(len(sample_image)):
+                    test_result = CocoPose.display_image(sample_image[i], heatMat[i], pafMat[i], as_numpy=True)
+                    test_result = cv2.resize(test_result, (640, 640))
+                    test_result = test_result.reshape([640, 640, 3]).astype(float)
+                    sample_results.append(test_result)
+
+                test_results = []
+                for i in range(len(val_image)):
+                    test_result = CocoPose.display_image(val_image[i], heatMat[len(sample_image) + i], pafMat[len(sample_image) + i], as_numpy=True)
+                    test_result = cv2.resize(test_result, (640, 640))
+                    test_result = test_result.reshape([640, 640, 3]).astype(float)
+                    test_results.append(test_result)
+
+                # save summary
+                summary = sess.run(merged_validate_op, feed_dict={
+                    valid_loss: average_loss / total_cnt,
+                    valid_loss_ll: average_loss_ll / total_cnt,
+                    valid_loss_ll_paf: average_loss_ll_paf / total_cnt,
+                    valid_loss_ll_heat: average_loss_ll_heat / total_cnt,
+                    sample_valid: test_results,
+                    sample_train: sample_results
+                })
+                file_writer.add_summary(summary, gs_num)
+
+        saver.save(sess, os.path.join(args.modelpath, training_name, 'model'), global_step=global_step)
+    logger.info('optimization finished. %f' % (time.time() - time_started))

posegait/src/webcam3d.py

@@ -0,0 +1,121 @@
+"""
+This serve as our base openGL class.
+"""
+
+import numpy as np
+import pyqtgraph.opengl as gl
+import pyqtgraph as pg
+from pyqtgraph.Qt import QtCore, QtGui
+import sys
+
+import cv2
+import time
+import os
+
+from estimator import TfPoseEstimator
+from networks import get_graph_path, model_wh
+from lifting.prob_model import Prob3dPose
+import common
+
+
+class Terrain(object):
+    def __init__(self):
+        """
+        Initialize the graphics window and mesh surface
+        """
+
+        # setup the view window
+        self.app = QtGui.QApplication(sys.argv)
+        self.window = gl.GLViewWidget()
+        self.window.setWindowTitle('Terrain')
+        self.window.setGeometry(0, 110, 1920, 1080)
+        self.window.setCameraPosition(distance=30, elevation=12)
+        self.window.show()
+
+        gx = gl.GLGridItem()
+        gy = gl.GLGridItem()
+        gz = gl.GLGridItem()
+        gx.rotate(90, 0, 1, 0)
+        gy.rotate(90, 1, 0, 0)
+        gx.translate(-10, 0, 0)
+        gy.translate(0, -10, 0)
+        gz.translate(0, 0, -10)
+        self.window.addItem(gx)
+        self.window.addItem(gy)
+        self.window.addItem(gz)
+
+        model = 'mobilenet_thin_432x368'
+        #model = 'mobilenet_thin'
+        camera = 0
+        w, h = model_wh(model)
+        self.e = TfPoseEstimator(get_graph_path(model), target_size=(w, h))
+        self.cam = cv2.VideoCapture(camera)
+        ret_val, image = self.cam.read()
+        self.poseLifting = Prob3dPose('./src/lifting/models/prob_model_params.mat')
+        keypoints = self.mesh(image)
+
+        self.points = gl.GLScatterPlotItem( 
+            pos=keypoints,
+            color=pg.glColor((0, 255, 0)),
+            size=15
+        )
+        self.window.addItem(self.points)
+
+    def mesh(self, image):
+        image_h, image_w = image.shape[:2]
+        width = 640
+        height = 480
+        pose_2d_mpiis = []
+        visibilities = []
+
+        humans = self.e.inference(image, scales=[None])
+
+        for human in humans:
+            pose_2d_mpii, visibility = common.MPIIPart.from_coco(human)
+            pose_2d_mpiis.append(
+                [(int(x * width + 0.5), int(y * height + 0.5)) for x, y in pose_2d_mpii]
+            )
+            visibilities.append(visibility)
+
+        pose_2d_mpiis = np.array(pose_2d_mpiis)
+        visibilities = np.array(visibilities)
+        transformed_pose2d, weights = self.poseLifting.transform_joints(pose_2d_mpiis, visibilities)
+        pose_3d = self.poseLifting.compute_3d(transformed_pose2d, weights)
+
+        keypoints = pose_3d[0].transpose()
+
+        return keypoints / 80
+
+    def update(self):
+        """
+        update the mesh and shift the noise each time
+        """
+        ret_val, image = self.cam.read()
+        try:
+            keypoints = self.mesh(image)
+        except AssertionError:
+            print('body not in image')
+        else:
+            self.points.setData(pos=keypoints)
+
+    def start(self):
+        """
+        get the graphics window open and setup
+        """
+        if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
+            QtGui.QApplication.instance().exec_()
+
+    def animation(self, frametime=10):
+        """
+        calls the update method to run in a loop
+        """
+        timer = QtCore.QTimer()
+        timer.timeout.connect(self.update)
+        timer.start(frametime)
+        self.start()
+
+
+if __name__ == '__main__':
+    os.chdir('..')
+    t = Terrain()
+    t.animation()