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, ...])