deeplabcut.pose_estimation_tensorflow.core.predict

Functions:

Name	Description
argmax_pose_predict	Combine scoremat and offsets to the final pose.
extract_cnn_output	Extract locref + scmap from network.
extract_cnn_outputmulti	extract locref + scmap from network
getpose	Extract pose.
getposeNP	Adapted from DeeperCut, performs numpy-based faster inference on batches.

argmax_pose_predict

argmax_pose_predict(scmap, offmat, stride)

Combine scoremat and offsets to the final pose.

Source code in deeplabcut/pose_estimation_tensorflow/core/predict.py

def argmax_pose_predict(scmap, offmat, stride):
    """Combine scoremat and offsets to the final pose."""
    num_joints = scmap.shape[2]
    pose = []
    for joint_idx in range(num_joints):
        maxloc = np.unravel_index(np.argmax(scmap[:, :, joint_idx]), scmap[:, :, joint_idx].shape)
        offset = np.array(offmat[maxloc][joint_idx])[::-1]
        pos_f8 = np.array(maxloc).astype("float") * stride + 0.5 * stride + offset
        pose.append(np.hstack((pos_f8[::-1], [scmap[maxloc][joint_idx]])))
    return np.array(pose)

extract_cnn_output

extract_cnn_output(outputs_np, cfg)

Extract locref + scmap from network.

Source code in deeplabcut/pose_estimation_tensorflow/core/predict.py

def extract_cnn_output(outputs_np, cfg):
    """Extract locref + scmap from network."""
    scmap = outputs_np[0]
    scmap = np.squeeze(scmap)
    locref = None
    if cfg["location_refinement"]:
        locref = np.squeeze(outputs_np[1])
        shape = locref.shape
        locref = np.reshape(locref, (shape[0], shape[1], -1, 2))
        locref *= cfg["locref_stdev"]
    if len(scmap.shape) == 2:  # for single body part!
        scmap = np.expand_dims(scmap, axis=2)
    return scmap, locref

extract_cnn_outputmulti

extract_cnn_outputmulti(outputs_np, cfg)

extract locref + scmap from network Dimensions: image batch x imagedim1 x imagedim2 x bodypart

Source code in deeplabcut/pose_estimation_tensorflow/core/predict.py

def extract_cnn_outputmulti(outputs_np, cfg):
    """extract locref + scmap from network
    Dimensions: image batch x imagedim1 x imagedim2 x bodypart"""
    scmap = outputs_np[0]
    locref = None
    if cfg["location_refinement"]:
        locref = outputs_np[1]
        shape = locref.shape
        locref = np.reshape(locref, (shape[0], shape[1], shape[2], -1, 2))
        locref *= cfg["locref_stdev"]
    if len(scmap.shape) == 2:  # for single body part!
        scmap = np.expand_dims(scmap, axis=2)
    return scmap, locref

getpose

getpose(image, cfg, sess, inputs, outputs, outall=False)

Extract pose.

Source code in deeplabcut/pose_estimation_tensorflow/core/predict.py

def getpose(image, cfg, sess, inputs, outputs, outall=False):
    """Extract pose."""
    im = np.expand_dims(image, axis=0).astype(float)
    outputs_np = sess.run(outputs, feed_dict={inputs: im})
    scmap, locref = extract_cnn_output(outputs_np, cfg)
    num_outputs = cfg.get("num_outputs", 1)
    if num_outputs > 1:
        pose = multi_pose_predict(scmap, locref, cfg["stride"], num_outputs)
    else:
        pose = argmax_pose_predict(scmap, locref, cfg["stride"])
    if outall:
        return scmap, locref, pose
    else:
        return pose

getposeNP

getposeNP(image, cfg, sess, inputs, outputs, outall=False)

Adapted from DeeperCut, performs numpy-based faster inference on batches.

Introduced in https://www.biorxiv.org/content/10.1101/457242v1

Source code in deeplabcut/pose_estimation_tensorflow/core/predict.py

def getposeNP(image, cfg, sess, inputs, outputs, outall=False):
    """Adapted from DeeperCut, performs numpy-based faster inference on batches.

    Introduced in https://www.biorxiv.org/content/10.1101/457242v1
    """

    num_outputs = cfg.get("num_outputs", 1)
    outputs_np = sess.run(outputs, feed_dict={inputs: image})

    scmap, locref = extract_cnn_outputmulti(outputs_np, cfg)  # processes image batch.
    batchsize, ny, nx, num_joints = scmap.shape

    Y, X = get_top_values(scmap, n_top=num_outputs)

    # Combine scoremat and offsets to the final pose.
    DZ = np.zeros((num_outputs, batchsize, num_joints, 3))
    for m in range(num_outputs):
        for l in range(batchsize):
            for k in range(num_joints):
                x = X[m, l, k]
                y = Y[m, l, k]
                DZ[m, l, k, :2] = locref[l, y, x, k, :]
                DZ[m, l, k, 2] = scmap[l, y, x, k]

    X = X.astype("float32") * cfg["stride"] + 0.5 * cfg["stride"] + DZ[:, :, :, 0]
    Y = Y.astype("float32") * cfg["stride"] + 0.5 * cfg["stride"] + DZ[:, :, :, 1]
    P = DZ[:, :, :, 2]

    Xs = X.swapaxes(0, 2).swapaxes(0, 1)
    Ys = Y.swapaxes(0, 2).swapaxes(0, 1)
    Ps = P.swapaxes(0, 2).swapaxes(0, 1)

    pose = np.empty((cfg["batch_size"], num_outputs * cfg["num_joints"] * 3), dtype=X.dtype)
    pose[:, 0::3] = Xs.reshape(batchsize, -1)
    pose[:, 1::3] = Ys.reshape(batchsize, -1)
    pose[:, 2::3] = Ps.reshape(batchsize, -1)

    if outall:
        return scmap, locref, pose
    else:
        return pose