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