File size: 996 Bytes
000b3e1 2e1b62d 000b3e1 2e1b62d 000b3e1 a5e6f78 000b3e1 2e1b62d 000b3e1 2e1b62d 000b3e1 2e1b62d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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
|