26 lines
622 B
Python
26 lines
622 B
Python
import scipy.io as sio # type: ignore
|
|
|
|
import torch
|
|
|
|
|
|
@torch.no_grad()
|
|
def make_mask(
|
|
filename_mask: str, data: torch.Tensor, device: torch.device, dtype: torch.dtype
|
|
) -> torch.Tensor:
|
|
mask: torch.Tensor = torch.tensor(
|
|
sio.loadmat(filename_mask)["maskInfo"]["maskIdx2D"][0][0],
|
|
device=device,
|
|
dtype=torch.bool,
|
|
)
|
|
mask = mask > 0.5
|
|
|
|
limit: torch.Tensor = torch.tensor(
|
|
2**16 - 1,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
if torch.any(data.flatten() >= limit):
|
|
mask = mask & ~(torch.any(torch.any(data >= limit, dim=-1), dim=-1))
|
|
|
|
return mask
|