preprocessors

dreem.datasets.preprocessors

Preprocessor classes for dataset preparation steps.

Classes:

Name	Description
NonMaxSuppression	Apply non-maximum suppression to a list of instances.
RemoveExcessDetections	Remove excess detections when the number exceeds max_tracks.

NonMaxSuppression

Bases: ProcessingStep

Apply non-maximum suppression to a list of instances.

This step applies non-maximum suppression to a list of instances based on their bounding box overlap.

Expected state keys

• instances: list of dreem.io.Instance objects
• max_detection_overlap: float - the maximum overlap allowed between instances

Methods:

Name	Description
__init__	Initialize NonMaxSuppression step.
run	Apply non-maximum suppression to a list of instances.

Source code in dreem/datasets/preprocessors.py

class NonMaxSuppression(ProcessingStep):
    """Apply non-maximum suppression to a list of instances.

    This step applies non-maximum suppression to a list of instances based on their bounding box overlap.

    Expected state keys:
        - instances: list of dreem.io.Instance objects
        - max_detection_overlap: float - the maximum overlap allowed between instances
    """

    def __init__(self, max_detection_overlap: float):
        """Initialize NonMaxSuppression step.

        Args:
            max_detection_overlap: The maximum overlap allowed between instances.
        """
        super().__init__(name="non_max_suppression")
        self.max_detection_overlap = max_detection_overlap

    def run(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Apply non-maximum suppression to a list of instances.

        Args:
            state: State dictionary with required keys.
                - frame_ind: int - index of the current frame
                - instances: list of dreem.io.Instance objects

        Returns:
            Modified state with instances removed.
        """
        instances = state["instances"]
        frame_ind = state["frame_ind"]
        discard = set()
        bboxes = np.stack([instance.bbox.squeeze(0) for instance in instances])
        ioms = pairwise_iom(Boxes(bboxes), Boxes(bboxes))
        high_iom_pairs = nms(ioms, self.max_detection_overlap)
        for pair in high_iom_pairs:
            if pair[0] in discard or pair[1] in discard:
                continue
            # Collect the 'pose' dictionary values for both instances and stack them into tensors
            inst_1_pose_tensor = torch.stack(
                [torch.tensor(v) for v in instances[pair[0]].pose.values()]
            )
            inst_2_pose_tensor = torch.stack(
                [torch.tensor(v) for v in instances[pair[1]].pose.values()]
            )
            inst_1_keypoints = len(
                inst_1_pose_tensor[~inst_1_pose_tensor.isnan().any(dim=1)]
            )
            inst_2_keypoints = len(
                inst_2_pose_tensor[~inst_2_pose_tensor.isnan().any(dim=1)]
            )
            # break ties by keeping the first instance
            if inst_1_keypoints >= inst_2_keypoints:
                id_to_discard = pair[1]
            else:
                id_to_discard = pair[0]
            discard.add(id_to_discard)
        for id in sorted(discard, reverse=True):
            removed = instances.pop(id)
            logger.warning(
                f"Removed instance with bounding box {removed.bbox.squeeze(0)} format [ymin, xmin, ymax, xmax] from frame {frame_ind} due to high bounding box overlap with another instance"
            )
        return state

__init__(max_detection_overlap)

Initialize NonMaxSuppression step.

Parameters:

Name	Type	Description	Default
max_detection_overlap	float	The maximum overlap allowed between instances.	required

Source code in dreem/datasets/preprocessors.py

def __init__(self, max_detection_overlap: float):
    """Initialize NonMaxSuppression step.

    Args:
        max_detection_overlap: The maximum overlap allowed between instances.
    """
    super().__init__(name="non_max_suppression")
    self.max_detection_overlap = max_detection_overlap

run(state)

Apply non-maximum suppression to a list of instances.

Parameters:

Name	Type	Description	Default
state	Dict[str, Any]	State dictionary with required keys. - frame_ind: int - index of the current frame - instances: list of dreem.io.Instance objects	required

Returns:

Type	Description
Dict[str, Any]	Modified state with instances removed.

Source code in dreem/datasets/preprocessors.py

def run(self, state: Dict[str, Any]) -> Dict[str, Any]:
    """Apply non-maximum suppression to a list of instances.

    Args:
        state: State dictionary with required keys.
            - frame_ind: int - index of the current frame
            - instances: list of dreem.io.Instance objects

    Returns:
        Modified state with instances removed.
    """
    instances = state["instances"]
    frame_ind = state["frame_ind"]
    discard = set()
    bboxes = np.stack([instance.bbox.squeeze(0) for instance in instances])
    ioms = pairwise_iom(Boxes(bboxes), Boxes(bboxes))
    high_iom_pairs = nms(ioms, self.max_detection_overlap)
    for pair in high_iom_pairs:
        if pair[0] in discard or pair[1] in discard:
            continue
        # Collect the 'pose' dictionary values for both instances and stack them into tensors
        inst_1_pose_tensor = torch.stack(
            [torch.tensor(v) for v in instances[pair[0]].pose.values()]
        )
        inst_2_pose_tensor = torch.stack(
            [torch.tensor(v) for v in instances[pair[1]].pose.values()]
        )
        inst_1_keypoints = len(
            inst_1_pose_tensor[~inst_1_pose_tensor.isnan().any(dim=1)]
        )
        inst_2_keypoints = len(
            inst_2_pose_tensor[~inst_2_pose_tensor.isnan().any(dim=1)]
        )
        # break ties by keeping the first instance
        if inst_1_keypoints >= inst_2_keypoints:
            id_to_discard = pair[1]
        else:
            id_to_discard = pair[0]
        discard.add(id_to_discard)
    for id in sorted(discard, reverse=True):
        removed = instances.pop(id)
        logger.warning(
            f"Removed instance with bounding box {removed.bbox.squeeze(0)} format [ymin, xmin, ymax, xmax] from frame {frame_ind} due to high bounding box overlap with another instance"
        )
    return state

RemoveExcessDetections

Bases: ProcessingStep

Remove excess detections when the number exceeds max_tracks.

This step mirrors the max_tracks pruning logic from dreem/datasets/sleap_dataset.py around line 487. It removes the lowest confidence instances when the total count exceeds the maximum allowed tracks.

Expected state keys

• instances: list of dreem.io.Instance objects

Modified state keys

• instances: list with excess low-confidence instances removed
• removed_detections: int - count of removed instances (if any were removed)

Methods:

Name	Description
__init__	Initialize RemoveExcessDetections step.
run	Remove excess detections based on instance scores.

Source code in dreem/datasets/preprocessors.py

class RemoveExcessDetections(ProcessingStep):
    """Remove excess detections when the number exceeds max_tracks.

    This step mirrors the max_tracks pruning logic from dreem/datasets/sleap_dataset.py
    around line 487. It removes the lowest confidence instances when the total count
    exceeds the maximum allowed tracks.

    Expected state keys:
        - instances: list of dreem.io.Instance objects

    Modified state keys:
        - instances: list with excess low-confidence instances removed
        - removed_detections: int - count of removed instances (if any were removed)
    """

    def __init__(self, max_tracks: int):
        """Initialize RemoveExcessDetections step.

        Args:
            max_tracks: The maximum number of tracks/detections allowed per frame.
        """
        super().__init__(name="remove_excess_detections")
        self.max_tracks = max_tracks

    def run(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Remove excess detections based on instance scores.

        Args:
            state: State dictionary with required keys.
                - frame_ind: int - index of the current frame
                - instances: list of dreem.io.Instance objects

        Returns:
            Modified state with excess instances removed.
        """
        frame_ind = state["frame_ind"]
        instances = state["instances"]

        removed = 0
        num_to_remove = len(instances) - self.max_tracks

        for _ in range(num_to_remove):
            lowest = min(instances, key=lambda x: x.instance_score)
            if lowest.instance_score < 1:
                instances.remove(lowest)
                removed += 1
        if removed > 0:
            logger.warning(
                f"Removed {removed} lowest confidence instances from frame {frame_ind} due to excess detections. Parameter 'max_tracks' in the tracker config sets the maximum number of detections per frame."
            )

        state["instances"] = instances
        state["removed_detections"] = removed

        return state

__init__(max_tracks)

Initialize RemoveExcessDetections step.

Parameters:

Name	Type	Description	Default
max_tracks	int	The maximum number of tracks/detections allowed per frame.	required

Source code in dreem/datasets/preprocessors.py

def __init__(self, max_tracks: int):
    """Initialize RemoveExcessDetections step.

    Args:
        max_tracks: The maximum number of tracks/detections allowed per frame.
    """
    super().__init__(name="remove_excess_detections")
    self.max_tracks = max_tracks

run(state)

Remove excess detections based on instance scores.

Parameters:

Name	Type	Description	Default
state	Dict[str, Any]	State dictionary with required keys. - frame_ind: int - index of the current frame - instances: list of dreem.io.Instance objects	required

Returns:

Type	Description
Dict[str, Any]	Modified state with excess instances removed.

Source code in dreem/datasets/preprocessors.py

def run(self, state: Dict[str, Any]) -> Dict[str, Any]:
    """Remove excess detections based on instance scores.

    Args:
        state: State dictionary with required keys.
            - frame_ind: int - index of the current frame
            - instances: list of dreem.io.Instance objects

    Returns:
        Modified state with excess instances removed.
    """
    frame_ind = state["frame_ind"]
    instances = state["instances"]

    removed = 0
    num_to_remove = len(instances) - self.max_tracks

    for _ in range(num_to_remove):
        lowest = min(instances, key=lambda x: x.instance_score)
        if lowest.instance_score < 1:
            instances.remove(lowest)
            removed += 1
    if removed > 0:
        logger.warning(
            f"Removed {removed} lowest confidence instances from frame {frame_ind} due to excess detections. Parameter 'max_tracks' in the tracker config sets the maximum number of detections per frame."
        )

    state["instances"] = instances
    state["removed_detections"] = removed

    return state