def process_batch(network, frames, box_centers, box_angles, box_size):
logging.info(f" exporting boxes.")
boxes, fly_id, fly_frame = export_boxes(frames,
box_centers,
box_size=box_size,
box_angles=box_angles)
logging.info(f" predicting confidence maps for {boxes.shape[0]} boxes.")
confmaps = predict_confmaps(network, normalize_boxes(boxes))
logging.info(f" processing confidence maps.")
positions, confidence = process_confmaps_simple(confmaps)
logging.info(f" recalculating box angles.")
# TODO: make positions self-documenting (named_list), make these args
head_idx = 0
tail_idx = 11
nb_flies = box_angles.shape[1]
newbox_angles, bad_boxes = detect_bad_boxes_by_angle(
unflatten(positions[:, head_idx, :], nb_flies),
unflatten(positions[:, tail_idx, :], nb_flies),
epsilon=5)
bad_frame_idx = np.any(bad_boxes,
axis=1)[:, 0] # for addressing bad boxes by frame
bad_box_idx = np.repeat(bad_frame_idx,
nb_flies) # for addressing bad boxes by box
fixed_angles = box_angles
if np.sum(bad_boxes) > 0:
logging.info(
f" found {np.sum(bad_boxes)} cases of boxes with angles above threshold."
)
logging.info(f" re-exporting the bad boxes.")
fixed_angles[bad_frame_idx,
...] = box_angles[bad_frame_idx,
...] + newbox_angles[bad_frame_idx, ...]
boxes[bad_box_idx, ...], *_ = export_boxes(
[frames[int(idx)] for idx in np.where(bad_frame_idx)[0]],
box_centers[bad_frame_idx, ...],
box_size=np.array([120, 120]),
box_angles=fixed_angles[bad_frame_idx, ...])
logging.info(f" re-doing predictions.")
confmaps[bad_box_idx, ...] = predict_confmaps(
network, normalize_boxes(boxes[bad_box_idx, ...]))
logging.info(f" re-processing confidence maps.")
positions[bad_box_idx, ...], confidence[bad_box_idx,
...] = process_confmaps_simple(
confmaps[bad_box_idx, ...])
# all results should be in nb_boxes format
fixed_angles = flatten(fixed_angles)
bad_boxes = flatten(bad_boxes)
return positions, confidence, confmaps, bad_boxes, fly_id, fly_frame, boxes, fixed_angles
def test_process_confmaps_simple():
X = np.zeros((4, 10, 10, 4))
X[0, 2, 4, 0] = 2.0
positions, confidence = process_confmaps_simple(X)
assert positions.shape == (4, 4, 2)
assert confidence.shape == (4, 4, 1)
assert confidence[0, 0, 0] == 1
def evaluate_network(network, boxes, positions, batch_size: int = 100):
"""Evaluate LEAP network on boxes and positions.
Args:
network
boxes: [nb_boxes, width, height, color-channels]
positions: [nb_boxes, nb_parts, x/y]
batch_size:32
Returns:
mean_map_error: MSE between predicted and actual confmaps
mean_position_error: MSE between prediction and actual positions
position_error: MSE between prediction and actual positions
"""
box_size = boxes.shape[1:3]
confmaps = make_masks(positions, size=box_size)
confmaps_predicted = predict_confmaps(network, boxes, batch_size)
mean_map_error = np.sum(np.square(confmaps - confmaps_predicted))
positions_predicted, confidence = process_confmaps_simple(confmaps)
position_error = np.square(positions - positions_predicted)
mean_position_error = np.mean(position_error)
return mean_map_error, mean_position_error, position_error
def process_batch(network, frames, box_centers, box_angles, box_size, priors):
logging.info(f" exporting boxes.")
boxes, fly_id, fly_frame = export_boxes(frames,
box_centers,
box_size=box_size,
box_angles=box_angles)
logging.info(f" predicting confidence maps for {boxes.shape[0]} boxes.")
confmaps = predict_confmaps(network, normalize_boxes(boxes))
logging.info(f" processing confidence maps.")
positions, confidence = process_confmaps_simple(confmaps)
logging.info(f" recalculating box angles.")
# TODO: make positions self-documenting (named_list), make these args
head_idx = 0
tail_idx = 11
nb_flies = box_angles.shape[1]
newbox_angles, bad_boxes = detect_bad_boxes_by_angle(
unflatten(positions[:, head_idx, :], nb_flies),
unflatten(positions[:, tail_idx, :], nb_flies),
epsilon=5)
bad_frame_idx = np.any(bad_boxes,
axis=1)[:, 0] # for addressing bad boxes by frame
bad_box_idx = np.repeat(bad_frame_idx,
nb_flies) # for addressing bad boxes by box
fixed_angles = box_angles
if np.sum(bad_boxes) > 0:
logging.info(
f" found {np.sum(bad_boxes)} cases of boxes with angles above threshold."
)
logging.info(f" re-exporting the bad boxes.")
fixed_angles[bad_frame_idx,
...] = box_angles[bad_frame_idx,
...] + newbox_angles[bad_frame_idx, ...]
boxes[bad_box_idx, ...], *_ = export_boxes(
[frames[int(idx)] for idx in np.where(bad_frame_idx)[0]],
box_centers[bad_frame_idx, ...],
box_size=np.array([120, 120]),
box_angles=fixed_angles[bad_frame_idx, ...])
logging.info(f" re-doing predictions.")
confmaps[bad_box_idx, ...] = predict_confmaps(
network, normalize_boxes(boxes[bad_box_idx, ...]))
logging.info(f" re-processing confidence maps.")
positions[bad_box_idx, ...], confidence[bad_box_idx,
...] = process_confmaps_simple(
confmaps[bad_box_idx, ...])
# all results should be in nb_boxes format
fixed_angles = flatten(fixed_angles)
bad_boxes = flatten(bad_boxes)
# Detection of errors for prior, according to position thresholds
error_matrix = detect_prior_cases(positions)
priorcase_idx = np.where(np.any(error_matrix, axis=1))[0]
logging.info(
f" number of boxes to process with priors: {priorcase_idx.shape[0]}")
# Try to fix positions applying priors
positions[priorcase_idx,
...] = priors_processing(positions[priorcase_idx, ...],
error_matrix[priorcase_idx, ...],
boxes[priorcase_idx,
...], confmaps[priorcase_idx,
...], priors)
# Re-evaluate error matrix
new_error_matrix = detect_prior_cases(positions)
new_idxs = np.where(np.any(new_error_matrix, axis=1))[0]
logging.info(
f" number of boxes still unfixed after priors: {new_idxs.shape[0]}")
return positions, confidence, confmaps, bad_boxes, fly_id, fly_frame, boxes, fixed_angles, np.any(
error_matrix, axis=1)
def main(expID: str = 'localhost-20180720_182837',
frame_start: int = 1000,
frame_stop: int = 2000,
frame_step: int = 100):
## Fix tracks
# Paths
trackPath = f"{resPath}\{expID}\{expID}_tracks.h5"
videoPath = f"{dataPath}\{expID}\{expID}.mp4"
trackfixedPath = f"{resPath}\{expID}\\{expID}_tracks_fixed.h5"
savingPath = f"{resPath}\{expID}"
posePath = f"{savingPath}\{expID}_pose.h5"
fixedBoxesPath = f"{savingPath}\{expID}_fixedBoxes.h5"
# Do not fix if they are already fixed
if not os.path.exists(trackfixedPath):
logging.info(f" doing adrian_fix_tracks")
adrian_fix_tracks(trackPath, trackfixedPath)
else:
logging.info(f" fixed tracks already exist")
## Load video
logging.info(f" loading video from {videoPath}.")
vr = VideoReader(videoPath)
## Load track
logging.info(f" loading tracks from {trackfixedPath}.")
has_tracks = False
try:
data = dd.io.load(trackfixedPath)
centers = data['centers'][:] # nframe, channel, fly id, coordinates
tracks = data['lines']
chbb = data['chambers_bounding_box'][:]
heads = tracks[:, 0, :, 0, ::-1] # nframe, fly id, coordinates
tails = tracks[:, 0, :, 1, ::-1] # nframe, fly id, coordinates
heads = heads + chbb[1][0][:] # nframe, fly id, coordinates
tails = tails + chbb[1][0][:] # nframe, fly id, coordinates
has_tracks = True
box_centers = centers[:, 0, :, :] # nframe, fly id, coordinates
box_centers = box_centers + chbb[1][0][:]
nb_flies = box_centers.shape[1]
logging.info(f" nflies: {nb_flies}.")
except OSError as e:
logging.error(f' could not load tracks.')
## Specifications for boxes
if frame_stop == 0:
frame_stop = data['frame_count']
logging.info(f' Setting frame_stop: {0}.'.format(frame_stop))
frame_range = range(frame_start, frame_stop, frame_step)
logging.info(f" frame range: {frame_start}:{frame_stop}:{frame_step}.")
## Create list of frames
logging.info(f" getting frames from video.")
frames = list(vr[frame_start:frame_stop:frame_step])
nb_frames = len(frames)
## Calculate angle for export boxes
logging.info(f" calculating box angles.")
box_angles = get_angles(heads[frame_range, ...], tails[frame_range, ...])
## Export boxes function
logging.info(f" exporting boxes of {expID}.")
boxes, fly_id, fly_frame = export_boxes(frames,
box_centers[frame_range, ...],
box_size=np.array([120, 120]),
box_angles=box_angles)
boxes = normalize_boxes(boxes)
# boxes = np.rot90(boxes,2,(1,2)) # Flips boxes, in case network expected fly images in the other direction
## Predictions
logging.info(f" doing predictions.")
confmaps = predict_confmaps(networkPath, boxes[:, :, :, :1])
logging.info(f" processing confidence maps.")
positions, confidence = process_confmaps_simple(confmaps)
## Recalculation of angles for further orientation fix
logging.info(f" recalculating box angles.")
newbox_angles, bad_boxes = detect_bad_boxes_byAngle(positions)
logging.info(
f" found {np.sum(bad_boxes)} cases of boxes with angles above threshold."
)
box_angles = box_angles + newbox_angles
logging.info(f" re-exporting boxes of {expID}.")
boxes, fly_id, fly_frame = export_boxes(frames,
box_centers[frame_range, ...],
box_size=np.array([120, 120]),
box_angles=box_angles)
boxes = normalize_boxes(boxes)
## Final predictions
logging.info(f" re-doing predictions.")
confmaps = predict_confmaps(networkPath, boxes[:, :, :, :1])
logging.info(f" re-processing confidence maps.")
positions, confidence = process_confmaps_simple(confmaps)
## Saving data
logging.info(f" saving pose predictions to: {posePath}.")
posedata = {
'positions': positions,
'confidence': confidence,
'confmaps': confmaps,
'expID': expID,
'frame_range': frame_range,
'fly_id': fly_id,
'fly_frame': fly_frame,
'bad_boxes': bad_boxes
}
dd.io.save(posePath, posedata)
logging.info(f" saving fixed boxes to: {fixedBoxesPath}.")
fixedBoxesdata = {'boxes': boxes, 'fly_id': fly_id, 'fly_frame': fly_frame}
dd.io.save(fixedBoxesPath, fixedBoxesdata)
## Play movie of the exported boxes
if play_boxes:
logging.info(f" playing video of boxes.")
vplay(boxes[fly_id == 0, ...])