# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_vision.io.ipynb.

# %% ../../nbs/00_vision.io.ipynb 1
from __future__ import annotations

# %% auto 0
__all__ = ['filter_masked', 'read_chn_file_as_tensor', 'read_multichan_files_as_tensor']

# %% ../../nbs/00_vision.io.ipynb 3
from PIL import Image
from fastai.vision.all import *

# %% ../../nbs/00_vision.io.ipynb 5
def filter_masked(raw_arr, in_msk: int, out_msk: int):
    "Replace input mask pixel value with selected value"
    return np.select([raw_arr == in_msk], [out_msk], raw_arr)

def read_chn_file_as_tensor(path: str) -> Tensor:
    "Read single channel file into tensor"
    img_arr = np.array(Image.open(path))
    msk_arr = filter_masked(img_arr, 55537, 9999)
    return Tensor(msk_arr / 10000)

def read_multichan_files_as_tensor(files: list(str)) -> Tensor:
    "Read individual channel tensor files into a tensor of channels"
    return torch.cat([read_chn_file_as_tensor(path)[None] for path in files])
