def load_config(filename = "pose_cfg.yaml"):
    if 'POSE_PARAM_PATH' in os.environ:
        filename = os.environ['POSE_PARAM_PATH'] + '/' + filename
    return cfg_from_file(filename) def create(cfg):
    dataset_type = cfg.dataset_type
    if dataset_type == "mpii":
        from dataset.mpii import MPII
        data = MPII(cfg)
    elif dataset_type == "coco":
        from dataset.mscoco import MSCOCO
        data = MSCOCO(cfg)
    elif dataset_type == "penn_action":
        from dataset.penn_action import PennAction
        data = PennAction(cfg)
    elif dataset_type == "default":
        data = PoseDataset(cfg)
    else:
        raise Exception("Unsupported dataset_type: \"{}\"".format(dataset_type))
    return data def setup_pose_prediction(cfg):
def extract_cnn_output(outputs_np, cfg, pairwise_stats = None):
def argmax_pose_predict(scmap, offmat, stride):
def argmax_arrows_predict(scmap, offmat, pairwise_diff, stride): def _npcircle(image, cx, cy, radius, color, transparency=0.0):
def check_point(cur_x, cur_y, minx, miny, maxx, maxy):
def visualize_joints(image, pose):
def show_heatmaps(cfg, img, scmap, pose, cmap="jet"):
def show_arrows(cfg, img, pose, arrows):
def waitforbuttonpress(): def data_to_input(data):
    return np.expand_dims(data, axis=0).astype(float) def extract_detections(cfg, scmap, locref, pairwise_diff):
    num_joints = cfg.num_joints
    num_pairwise_relations = pairwise_diff.shape[2]

    # get dist_grid
    dist_grid = make_nms_grid(cfg.nms_radius)

    unProb = [None] * num_joints
    unPos = [None] * num_joints
    unPos_grid = [None] * num_joints
    pairwiseDiff = [None] * num_joints

    # apply nms
    for p_idx in range(num_joints):
        # IMPORTANT, as C++ function expects row-major
        prob_map = np.ascontiguousarray(scmap[:, :, p_idx])
        # print(prob_map.flags) has to be C_CONTIGUOUS

        dets = nms_grid(prob_map, dist_grid, cfg.det_min_score)

        cur_prob = np.zeros([len(dets), 1], dtype=np.float64)
        cur_pos = np.zeros([len(dets), 2], dtype=np.float64)
        cur_pos_grid = np.zeros([len(dets), 2], dtype=np.float64)
        cur_pairwise = np.zeros([len(dets), num_pairwise_relations, 2], dtype=np.float64)

        for idx, didx in enumerate(dets):
            ix = didx % scmap.shape[1]
            iy = didx // scmap.shape[1]

            cur_prob[idx, 0] = scmap[iy, ix, p_idx]
            cur_pos_grid[idx, :] = pos_from_grid_raw(cfg, np.array([ix, iy]))
            cur_pos[idx, :] = cur_pos_grid[idx, :] + locref[iy, ix, p_idx, :]
            cur_pairwise[idx, :, :] = pairwise_diff[iy, ix, :, :]

        unProb[p_idx] = cur_prob
        unPos[p_idx] = cur_pos
        unPos_grid[p_idx] = cur_pos_grid
        pairwiseDiff[p_idx] = cur_pairwise

    return Detections(coord=unPos, coord_grid=unPos_grid, conf=unProb, pairwise=pairwiseDiff) class SpatialModel:

    def __init__(self, cfg):
        self.num_keypoints = cfg.num_joints
        num_keypoints = cfg.num_joints
        self.cfg = cfg

        self.graph_dict = dict()

        self.same_part_pw_coef = 0.2

        self.X_min = [[None]*num_keypoints for idx in range(num_keypoints)]
        self.X_max = [[None]*num_keypoints for idx in range(num_keypoints)]
        self.w = [[None]*num_keypoints for idx in range(num_keypoints)]

    def load(self):
        for cidx1 in range(self.num_keypoints):
            for cidx2 in range(cidx1+1, self.num_keypoints):
                model_name  = "{}/spatial_model_cidx_{}_{}.mat".format(self.cfg.pairwise_model_dir, cidx1 + 1, cidx2 + 1)
                #print "loading:", model_name
                if not os.path.isfile(model_name):
                    continue
    
                spatial_model = sio.loadmat(model_name)

                self.X_max[cidx1][cidx2] = spatial_model['spatial_model']['training_opts'][0][0][0]['X_max'][0][0]
                self.X_min[cidx1][cidx2] = spatial_model['spatial_model']['training_opts'][0][0][0]['X_min'][0][0]
                self.w[cidx1][cidx2] = spatial_model['spatial_model']['log_reg'][0][0][0]['w'][0][0][:]

                #self.pair_class[cidx1][cidx2] = PairClassDistAngle(w=w, X_min=X_min, X_max=X_max)
                
        if not self.cfg.tensorflow_pairwise_order:
            tmpval = sio.loadmat(self.cfg.pairwise_stats_fn)
            self.graph = tmpval['graph']
            self.pairwise_means = tmpval['means']
            self.pairwise_std_devs = tmpval['std_devs']

            for gidx in range(self.graph.shape[0]):
                cidx1 = self.graph[gidx, 0]
                cidx2 = self.graph[gidx, 1]
                self.graph_dict[(cidx1, cidx2)] = gidx

    def get_fwd_bwd_index(self, cidx1, cidx2):
        if self.cfg.tensorflow_pairwise_order:
            fwd_idx = get_pairwise_index(cidx1, cidx2, self.cfg.num_joints)
            bwd_idx = get_pairwise_index(cidx2, cidx1, self.cfg.num_joints)
        else:
            fwd_idx = self.graph_dict[(cidx1+1, cidx2+1)]
            bwd_idx = self.graph_dict[(cidx2+1, cidx1+1)]
        return fwd_idx, bwd_idx

    def need_this_pairwise(self, cidx1, cidx2):
        if cidx1 == cidx2:
            return True
        sparse_graph = self.cfg.sparse_graph
        return not sparse_graph or [cidx1, cidx2] in sparse_graph

    def eval(self, cidx1, cidx2, detections):
        unPos = detections.coord

        idx_type1 = np.array(range(unPos[cidx1].shape[0]))
        idx_type2 = np.array(range(unPos[cidx2].shape[0]))

        assert(idx_type1.shape[0] > 0)
        assert(idx_type2.shape[0] > 0)

        num_edges = len(idx_type1) * len(idx_type2)

        tmpidx1, tmpidx2 = np.meshgrid(idx_type1, idx_type2)
        ptidx = np.hstack((tmpidx1.T.reshape((num_edges, 1)), tmpidx2.T.reshape((num_edges, 1))))

        if cidx1 != cidx2:
            cur_prob = self.compute_different_part_pairwise(cidx1, cidx2, detections, ptidx, num_edges)
        else:
            cur_prob = None
            ptidx = ptidx[ptidx[:, 0] < ptidx[:, 1]]

            delta = unPos[cidx2][ptidx[:, 1], :] - unPos[cidx1][ptidx[:, 0], :]
            dists = np.linalg.norm(delta, axis=1, ord=2)

            cur_prob = 1./(1 + np.exp(self.same_part_pw_coef*dists-7.5))

        return cur_prob, ptidx

    def compute_different_part_pairwise(self, cidx1, cidx2, detections, ptidx, num_edges):
        unPos = detections.coord
        unPos_grid = detections.coord_grid
        nextReg = detections.pairwise

        fwd_idx, bwd_idx = self.get_fwd_bwd_index(cidx1, cidx2)
        #print("cidxs ", cidx1, cidx2)
        #print("pairwise index", fwd_idx, bwd_idx)

        assert (ptidx.shape[0] > 0)

        delta_real_forward = unPos[cidx2][ptidx[:, 1], :] - unPos_grid[cidx1][ptidx[:, 0], :]
        delta_real_backward = unPos[cidx1][ptidx[:, 0], :] - unPos_grid[cidx2][ptidx[:, 1], :]

        delta_forward = nextReg[cidx1][ptidx[:, 0], fwd_idx, :]
        delta_backward = nextReg[cidx2][ptidx[:, 1], bwd_idx, :]

        dist1, abs_a1 = compute_features(delta_real_forward, delta_forward)
        dist2, abs_a2 = compute_features(delta_real_backward, delta_backward)

        featAugm = np.hstack((dist1.reshape(num_edges, 1), abs_a1.reshape(num_edges, 1), dist2.reshape(num_edges, 1),
                              abs_a2.reshape(num_edges, 1)))

        # print "features: ", cidx1, cidx2, featAugm

        featAugm = np.hstack((featAugm, np.exp(-featAugm), np.ones((num_edges, 1))))

        featAugm[:, :-1] = (featAugm[:, :-1] - self.X_min[cidx1][cidx2]) / (
            self.X_max[cidx1][cidx2] - self.X_min[cidx1][cidx2])
        cur_prob = 1 / (1 + np.exp(-featAugm.dot(self.w[cidx1][cidx2])))

        return cur_prob def eval_graph(sm, detections):

    time_start = time.time()

    unary_prob = detections.conf
    coordinates = detections.coord
    pairwise_regr = detections.pairwise

    cidx_list = range(0, sm.num_keypoints)
    #cidx_list = [0, 1, 2]

    unary_counts = []

    for idx, cidx in enumerate(cidx_list):
        unary_counts.append(unary_prob[cidx].shape[0])
    
    num_unary = sum(unary_counts)
    num_pairwise = 0

    for idx1, cidx1 in enumerate(cidx_list):
        for cidx2 in cidx_list[idx1:]:
            if cidx1 == cidx2:
                num_pairwise += unary_prob[cidx1].shape[0]*(unary_prob[cidx1].shape[0] - 1) // 2
            else:
                num_pairwise += unary_prob[cidx1].shape[0]*unary_prob[cidx2].shape[0]

    pos_array = np.zeros([num_unary, 2], dtype=np.float64)
    unary_array = np.zeros([num_unary, 1], dtype=np.float64)
    pw_array = np.zeros([num_pairwise, 1], dtype=np.float64)
    pwidx_array = np.zeros([num_pairwise, 2], dtype=np.uint16)

    firstidx = 0
    firstidx_list = []
    for idx, cidx in enumerate(cidx_list):
        lastidx = firstidx + unary_prob[cidx].shape[0]
        unary_array[firstidx:lastidx] = unary_prob[cidx]
        pos_array[firstidx:lastidx] = coordinates[cidx]

        firstidx_list.append(firstidx)
        firstidx = lastidx

    firstidx = 0
    for idx1, cidx1 in enumerate(cidx_list):
        for idx2 in range(idx1, len(cidx_list)):

            if coordinates[cidx1].shape[0] > 0:
                cidx2 = cidx_list[idx2]

                if coordinates[cidx2].shape[0] > 0:

                    if not sm.need_this_pairwise(cidx1, cidx2):
                        continue

                    cur_prob, ptidx = sm.eval(cidx1, cidx2, detections)
                    lastidx = firstidx + cur_prob.shape[0]

                    ptidx[:, 0] += firstidx_list[idx1]
                    ptidx[:, 1] += firstidx_list[idx2]

                    pw_array[firstidx:lastidx, 0] = cur_prob
                    pwidx_array[firstidx:lastidx, :] = ptidx

                    firstidx = lastidx

    is_sparse_graph = True
    solver_type = False
    do_suppression = True
    logit_in_solver = False

    if unary_array.shape[0] > 0:

        unary_array_solver = unary_array if logit_in_solver else logit_transform(unary_array)
        pw_array_solver = pw_array if logit_in_solver else logit_transform(pw_array)

        time_start = time.time()

        res = solve_nl_lmp(unary_array_solver, pwidx_array, pw_array_solver,
                          is_sparse_graph, solver_type, do_suppression, logit_in_solver)

        unLab = np.array(res, dtype=np.uint64)

        firstidx = 0
        for cidx in cidx_list:
            lastidx = firstidx + unary_prob[cidx].shape[0]
            unLab[firstidx:lastidx, 0] = cidx
            firstidx = lastidx

    else:
        unLab = np.array([])

    return unLab, pos_array, unary_array, pwidx_array, pw_array def get_person_conf_multicut(sm, unLab, unary_array, pos_array):
    if unLab.shape[0] > 0:
        num_people = int(np.max(unLab[:, 1])) + 1
    else:
        num_people = 0

    person_conf = np.zeros([num_people, sm.num_keypoints, 2])
    sum_prob = np.zeros([num_people, sm.num_keypoints, 1])

    # compute weighted average of keypoints of the same time 
    for didx in range(unLab.shape[0]):
        kidx = unLab[didx,0]
        pidx = unLab[didx,1]

        person_conf[pidx, kidx, :] += pos_array[didx, :]*unary_array[didx]
        sum_prob[pidx, kidx] += unary_array[didx]

    print("num_people: ", num_people)

    for pidx in range(num_people):
        for kidx in range(sm.num_keypoints):
            if sum_prob[pidx, kidx] > 0:
                person_conf[pidx, kidx, :] /= sum_prob[pidx, kidx]

    return person_conf class PersonDraw:
    def __init__(self):
        self.mk = munkres.Munkres()

        self.prev_person_conf = np.zeros([0, 1])
        self.prev_color_assignment = None

        # generated colors from http://tools.medialab.sciences-po.fr/iwanthue/
        track_colors_str = ["#F5591E",
                            "#3870FB",
                            "#FE5DB0",
                            "#B4A691",
                            "#43053F",
                            "#3475B1",
                            "#642612",
                            "#B3B43D",
                            "#DD9BFE",
                            "#28948D",
                            "#E99D53",
                            "#012B46",
                            "#9D2DA3",
                            "#04220A",
                            "#62CB22",
                            "#EE8F91",
                            "#D71638",
                            "#00613A",
                            "#318918",
                            "#B770FF",
                            "#82C091",
                            "#6C1333",
                            "#973405",
                            "#B19CB2",
                            "#F6267B",
                            "#284489",
                            "#97BF17",
                            "#3B899C",
                            "#931813",
                            "#FA76B6"]

        self.track_colors = [(int(s[1:3], 16), int(s[3:5], 16), int(s[5:7], 16)) for s in track_colors_str]

    def draw(self, visim, dataset, person_conf):
        minx = 2 * marker_size
        miny = 2 * marker_size
        maxx = visim.shape[1] - 2 * marker_size
        maxy = visim.shape[0] - 2 * marker_size

        num_people = person_conf.shape[0]
        color_assignment = dict()

        # MA: assign same color to matching body configurations
        if self.prev_person_conf.shape[0] > 0 and person_conf.shape[0] > 0:
            ref_points = get_ref_points(person_conf)
            prev_ref_points = get_ref_points(self.prev_person_conf)

            # MA: this munkres implementation assumes that num(rows) >= num(columns)
            if person_conf.shape[0] <= self.prev_person_conf.shape[0]:
                cost_matrix = scipy.spatial.distance.cdist(ref_points, prev_ref_points)
            else:
                cost_matrix = scipy.spatial.distance.cdist(prev_ref_points, ref_points)

            assert (cost_matrix.shape[0] <= cost_matrix.shape[1])

            conf_assign = self.mk.compute(cost_matrix)

            if person_conf.shape[0] > self.prev_person_conf.shape[0]:
                conf_assign = [(idx2, idx1) for idx1, idx2 in conf_assign]
                cost_matrix = cost_matrix.T

            for pidx1, pidx2 in conf_assign:
                if cost_matrix[pidx1][pidx2] < min_match_dist:
                    color_assignment[pidx1] = self.prev_color_assignment[pidx2]

        print("#tracked objects:", len(color_assignment))

        free_coloridx = sorted(list(set(range(len(self.track_colors))).difference(set(color_assignment.values()))),
                               reverse=True)

        for pidx in range(num_people):
            # color_idx = pidx % len(self.track_colors)
            if pidx in color_assignment:
                color_idx = color_assignment[pidx]
            else:
                if len(free_coloridx) > 0:
                    color_idx = free_coloridx[-1]
                    free_coloridx = free_coloridx[:-1]
                else:
                    color_idx = np.random.randint(len(self.track_colors))

                color_assignment[pidx] = color_idx

            assert (color_idx < len(self.track_colors))

            if np.sum(person_conf[pidx, :, 0] > 0) < draw_conf_min_count:
                continue

            for kidx1, kidx2 in dataset.get_pose_segments():
                p1 = (int(math.floor(person_conf[pidx, kidx1, 0])), int(math.floor(person_conf[pidx, kidx1, 1])))
                p2 = (int(math.floor(person_conf[pidx, kidx2, 0])), int(math.floor(person_conf[pidx, kidx2, 1])))

                if check_point(p1[0], p1[1], minx, miny, maxx, maxy) and check_point(p2[0], p2[1], minx, miny, maxx,
                                                                                     maxy):
                    color = np.array(self.track_colors[color_idx][::-1], dtype=np.float64) / 255.0
                    plt.plot([p1[0], p2[0]], [p1[1], p2[1]], marker='o', linestyle='solid', linewidth=2.0, color=color)


        self.prev_person_conf = person_conf
        self.prev_color_assignment = color_assignment def visualize_detections(cfg, img, detections):
    vis_scale = 1.0
    marker_size = 4

    minx = 2 * marker_size
    miny = 2 * marker_size
    maxx = img.shape[1] - 2 * marker_size
    maxy = img.shape[0] - 2 * marker_size

    unPos = detections.coord
    joints_to_visualise = range(cfg.num_joints)
    visim_dets = img.copy()
    for pidx in joints_to_visualise:
        for didx in range(unPos[pidx].shape[0]):
            cur_x = unPos[pidx][didx, 0] * vis_scale
            cur_y = unPos[pidx][didx, 1] * vis_scale

            # / cfg.global_scale

            if check_point(cur_x, cur_y, minx, miny, maxx, maxy):
                _npcircle(visim_dets,
                          cur_x, cur_y,
                          marker_size,
                          keypoint_colors[pidx])
    return visim_dets
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.misc import imread, imsave
from util.config import load_config
from dataset.factory import create as create_dataset
from nnet import predict
from util import visualize
from dataset.pose_dataset import data_to_input
from multiperson.detections import extract_detections
from multiperson.predict import SpatialModel, eval_graph, get_person_conf_multicut
from multiperson.visualize import PersonDraw, visualize_detections


sys.path.append(os.path.dirname(__file__) + "/../")






cfg = load_config("demo/pose_cfg_multi.yaml")

dataset = create_dataset(cfg)

sm = SpatialModel(cfg)
sm.load()

draw_multi = PersonDraw()

# Load and setup CNN part detector

sess, inputs, outputs = predict.setup_pose_prediction(cfg)