import torch


def get_box_from_mask(mask: torch.Tensor) -> torch.Tensor:
    """
    Returns the bounding box for a particular segmentation mask
    :param: mask torch.Tensor[X,Y] some mask where 0 is background and !0 is a segmentation mask
    :return: torch.Tensor[4] coordinates of the box surrounding the segmentation mask [x1, y1, x2, y2]
    """
    ind = torch.nonzero(mask)

    if ind.shape[0] == 0:
        box = torch.tensor([0, 0, 0, 0]).to(mask.device)

    else:
        box = torch.empty(4).to(mask.device)
        x = ind[:, 1]
        y = ind[:, 0]
        torch.stack((torch.min(x), torch.min(y), torch.max(x), torch.max(y)), out=box)

    return box