import pickle


def load_detectron_predictions(detectron_root):
    """Load detectron predictions from root directort.

    Args:
        detectron_root (Path): Points to a directory which contains a
            subdirectory for each sequence, which in turn contains a .pickle
            file for each frame in the sequence.

    Returns:
        predictions (dict): Map sequence to dict mapping from frame name to
            a dictionary containining keys 'boxes', and 'segmentations'.
    """
    predictions = {}
    for sequence_path in detectron_root.iterdir():
        if not sequence_path.is_dir():
            continue

        sequence = sequence_path.stem
        predictions[sequence] = {}
        for detectron_path in sequence_path.glob('*.pickle'):
            with open(detectron_path, 'rb') as f:
                frame_data = pickle.load(f)

            if frame_data['segmentations'] is None:
                frame_data['segmentations'] = [
                    [] for _ in range(len(frame_data['boxes']))
                ]

            frame_name = detectron_path.stem
            predictions[sequence][frame_name] = {
                'boxes': [],
                'segmentations': []
            }
            for c in range(len(frame_data['boxes'])):
                predictions[sequence][frame_name]['boxes'].extend(
                    frame_data['boxes'][c])
                predictions[sequence][frame_name]['segmentations'].extend(
                    frame_data['segmentations'][c])
    return predictions