select by roi aggregation

select.py

import numpy as np
from dipy.tracking.utils import near_roi

def reduce_rois(rois, include):
    include_roi = np.zeros(rois[0].shape, dtype=bool)
    exclude_roi = np.zeros(rois[0].shape, dtype=bool)

    for i in range(len(rois)):
        if include[i]:
            include_roi |= rois[i]
        else:
            exclude_roi |= rois[i]

    return include_roi, exclude_roi


def select(streamlines, include_roi, exclude_roi, affine, tol):
    include = near_roi(streamlines, include_roi, affine, tol)
    exclude = near_roi(streamlines, exclude_roi, affine, tol)
    return include & (~exclude)


def select_by_roi(streamlines, rois, include, affine=None, tol=0):
    include_roi, exclude_roi = reduce_rois(rois, include)
    return select(streamlines, include_roi, exclude_roi, affine, tol)