| 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 | |