def GetPoseS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes):
''' Non batch wise pose estimation for video cap.'''
if cfg['cropping']:
ny, nx = checkcropping(cfg, cap)
PredictedData = np.zeros(
(nframes,
dlc_cfg['num_outputs'] * 3 * len(dlc_cfg['all_joints_names'])))
pbar = tqdm(total=nframes)
counter = 0
step = max(10, int(nframes / 100))
while (cap.isOpened()):
if counter % step == 0:
pbar.update(step)
ret, frame = cap.read()
if ret:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
if cfg['cropping']:
frame = img_as_ubyte(frame[cfg['y1']:cfg['y2'],
cfg['x1']:cfg['x2']])
else:
frame = img_as_ubyte(frame)
pose = predict.getpose(frame, dlc_cfg, sess, inputs, outputs)
PredictedData[counter, :] = pose.flatten(
) # NOTE: thereby cfg['all_joints_names'] should be same order as bodyparts!
else:
nframes = counter
break
counter += 1
pbar.close()
return PredictedData, nframes
def get_pose(image, config, session, inputs, outputs):
"""
Gets scoremap, local reference and pose from DeepLabCut using given image
Pose is most probable points for each joint, and not really used later
Scoremap and local reference is essential to extract skeletons
:param image: frame which would be analyzed
:param config, session, inputs, outputs: DeepLabCut configuration and TensorFlow variables from load_deeplabcut()
:return: tuple of scoremap, local reference and pose
"""
scmap, locref, pose = predict.getpose(image, config, session, inputs, outputs, True)
return scmap, locref, pose
def GetPosesofFrames(cfg, dlc_cfg, sess, inputs, outputs, directory, framelist,
nframes, batchsize, rgb):
''' Batchwise prediction of pose for frame list in directory'''
#from skimage.io import imread
from deeplabcutcore.utils.auxfun_videos import imread
print("Starting to extract posture")
if rgb:
im = imread(os.path.join(directory, framelist[0]), mode='RGB')
else:
im = imread(os.path.join(directory, framelist[0]))
ny, nx, nc = np.shape(im)
print("Overall # of frames: ", nframes,
" found with (before cropping) frame dimensions: ", nx, ny)
PredictedData = np.zeros(
(nframes,
dlc_cfg['num_outputs'] * 3 * len(dlc_cfg['all_joints_names'])))
batch_ind = 0 # keeps track of which image within a batch should be written to
batch_num = 0 # keeps track of which batch you are at
if cfg['cropping']:
print(
"Cropping based on the x1 = %s x2 = %s y1 = %s y2 = %s. You can adjust the cropping coordinates in the config.yaml file."
% (cfg['x1'], cfg['x2'], cfg['y1'], cfg['y2']))
nx, ny = cfg['x2'] - cfg['x1'], cfg['y2'] - cfg['y1']
if nx > 0 and ny > 0:
pass
else:
raise Exception('Please check the order of cropping parameter!')
if cfg['x1'] >= 0 and cfg['x2'] < int(
np.shape(im)[1]) and cfg['y1'] >= 0 and cfg['y2'] < int(
np.shape(im)[0]):
pass #good cropping box
else:
raise Exception('Please check the boundary of cropping!')
pbar = tqdm(total=nframes)
counter = 0
step = max(10, int(nframes / 100))
if batchsize == 1:
for counter, framename in enumerate(framelist):
#frame=imread(os.path.join(directory,framename),mode='RGB')
if rgb:
im = imread(os.path.join(directory, framename), mode='RGB')
else:
im = imread(os.path.join(directory, framename))
if counter % step == 0:
pbar.update(step)
if cfg['cropping']:
frame = img_as_ubyte(im[cfg['y1']:cfg['y2'],
cfg['x1']:cfg['x2'], :])
else:
frame = img_as_ubyte(im)
pose = predict.getpose(frame, dlc_cfg, sess, inputs, outputs)
PredictedData[counter, :] = pose.flatten()
else:
frames = np.empty(
(batchsize, ny, nx, 3),
dtype='ubyte') # this keeps all the frames of a batch
for counter, framename in enumerate(framelist):
if rgb:
im = imread(os.path.join(directory, framename), mode='RGB')
else:
im = imread(os.path.join(directory, framename))
if counter % step == 0:
pbar.update(step)
if cfg['cropping']:
frames[batch_ind] = img_as_ubyte(im[cfg['y1']:cfg['y2'],
cfg['x1']:cfg['x2'], :])
else:
frames[batch_ind] = img_as_ubyte(im)
if batch_ind == batchsize - 1:
pose = predict.getposeNP(frames, dlc_cfg, sess, inputs,
outputs)
PredictedData[batch_num * batchsize:(batch_num + 1) *
batchsize, :] = pose
batch_ind = 0
batch_num += 1
else:
batch_ind += 1
if batch_ind > 0: #take care of the last frames (the batch that might have been processed)
pose = predict.getposeNP(
frames, dlc_cfg, sess, inputs, outputs
) #process the whole batch (some frames might be from previous batch!)
PredictedData[batch_num * batchsize:batch_num * batchsize +
batch_ind, :] = pose[:batch_ind, :]
pbar.close()
return PredictedData, nframes, nx, ny
def GetPoseDynamic(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes,
detectiontreshold, margin):
''' Non batch wise pose estimation for video cap by dynamically cropping around previously detected parts.'''
if cfg['cropping']:
ny, nx = checkcropping(cfg, cap)
else:
ny, nx = (int(cap.get(4)), int(cap.get(3)))
x1, x2, y1, y2 = 0, nx, 0, ny
detected = False
#TODO: perform detection on resized image (For speed)
PredictedData = np.zeros((nframes, 3 * len(dlc_cfg['all_joints_names'])))
pbar = tqdm(total=nframes)
counter = 0
step = max(10, int(nframes / 100))
while (cap.isOpened()):
if counter % step == 0:
pbar.update(step)
ret, frame = cap.read()
if ret:
#print(counter,x1,x2,y1,y2,detected)
originalframe = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
if cfg['cropping']:
frame = img_as_ubyte(originalframe[cfg['y1']:cfg['y2'],
cfg['x1']:cfg['x2']])[y1:y2,
x1:x2]
else:
frame = img_as_ubyte(originalframe[y1:y2, x1:x2])
pose = predict.getpose(frame, dlc_cfg, sess, inputs,
outputs).flatten()
detection = np.any(
pose[2::3] > detectiontreshold) #is anything detected?
if detection:
pose[0::3], pose[1::3] = pose[0::3] + x1, pose[
1::3] + y1 #offset according to last bounding box
x1, x2, y1, y2 = getboundingbox(
pose[0::3], pose[1::3], nx, ny,
margin) #coordinates for next iteration
if not detected:
detected = True #object detected
else:
if detected and (
x1 + y1 + y2 - ny + x2 - nx
) != 0: #was detected in last frame and dyn. cropping was performed >> but object lost in cropped variant >> re-run on full frame!
#print("looking again, lost!")
if cfg['cropping']:
frame = img_as_ubyte(
originalframe[cfg['y1']:cfg['y2'],
cfg['x1']:cfg['x2']])
else:
frame = img_as_ubyte(originalframe)
pose = predict.getpose(
frame, dlc_cfg, sess, inputs,
outputs).flatten() #no offset is necessary
x0, y0 = x1, y1
x1, x2, y1, y2 = 0, nx, 0, ny
detected = False
PredictedData[counter, :] = pose
else:
nframes = counter
break
counter += 1
pbar.close()
return PredictedData, nframes
