dlc_utils.py

import numpy as np
from dlclive import DLCLive


##########################################
def predict_dlc(list_np_crops, kpts_likelihood_th, dlc_model_folder, dlc_proc):

    # no animal detected: nothing to run
    if len(list_np_crops) == 0:
        return []

    # DLCLive always converts 3-channel frames BGR->RGB, but our crops are
    # already RGB (PIL), so pass them as BGR for the model to see RGB
    list_np_crops = [np.ascontiguousarray(crop[..., ::-1]) for crop in list_np_crops]

    # run dlc thru list of crops
    dlc_live = DLCLive(dlc_model_folder, processor=dlc_proc)
    dlc_live.init_inference(list_np_crops[0])

    list_kpts_per_crop = []
    for crop in list_np_crops:
        keypts_xyp = dlc_live.get_pose(crop)  # third column is llk!
        # set kpts below threhsold to nan
        keypts_xyp[keypts_xyp[:, -1] < kpts_likelihood_th, :] = np.nan
        # add kpts of this crop to list
        list_kpts_per_crop.append(keypts_xyp)

    return list_kpts_per_crop