Skip to content

collate

data.collate

default_device module-attribute

default_device = device('cpu')

images_to_4d_tensor

images_to_4d_tensor(images: List[Union[ndarray, Image, List]], device: device = default_device, apply_contiguous_div255: bool = True) -> Tensor

Converts a list of images to a 4D tensor.

Parameters:

  • images (List[Union[ndarray, Image, List]]) –

    images to convert.

  • device (device, default: default_device ) –

    device to move the tensor to. Defaults to cpu.

  • apply_contiguous_div255 (bool, default: True ) –

    whether to apply contiguous_div255. Defaults to True.

Returns:

  • Tensor

    torch.Tensor: 4D tensor of images.

Source code in SaigeToolkit/data/collate.py
def images_to_4d_tensor(
    images: List[Union[np.ndarray, Image.Image, List]],
    device: torch.device = default_device,
    apply_contiguous_div255: bool = True,
) -> torch.Tensor:
    """Converts a list of images to a 4D tensor.

    Args:
        images (List[Union[np.ndarray, Image.Image, List]]): images to convert.
        device (torch.device, optional): device to move the tensor to. Defaults to cpu.
        apply_contiguous_div255 (bool, optional): whether to apply contiguous_div255. Defaults to True.

    Returns:
        torch.Tensor: 4D tensor of images.

    """

    # multipage인 경우 List[image] -> n channel image로 변환
    if isinstance(images[0], List):
        images = [_merge_multipage_to_n_channel_array(image) for image in images]

    # list -> np.ndarray (BHWC) -> BCHW
    if isinstance(images[0], np.ndarray) and len(images) == 1:
        data = images[0][np.newaxis]
    else:
        data = np.stack(images)

    if data.ndim == 3:  # gray images
        data = data[:, None]
    else:
        data = data.transpose(0, 3, 1, 2)

    # array -> tensor
    data = torch.from_numpy(data).to(device)

    if apply_contiguous_div255:
        data = _contiguous_div255(tensor=data)

    return data

_contiguous_div255

_contiguous_div255(tensor: Tensor)
Source code in SaigeToolkit/data/collate.py
def _contiguous_div255(tensor: torch.Tensor):
    # contiguous -> uint8 -> float -> [0, 1] normalize
    return tensor.contiguous().div(255)

_merge_multipage_to_n_channel_array

_merge_multipage_to_n_channel_array(image: List[Union[ndarray, Image]]) -> ndarray
Source code in SaigeToolkit/data/collate.py
def _merge_multipage_to_n_channel_array(image: List[Union[np.ndarray, Image.Image]]) -> np.ndarray:
    # k장의 이미지를 n채널 이미지(H x W x n array)로 concat합니다.
    merged = []
    for image_i in image:
        if isinstance(image_i, Image.Image):
            image_i = np.array(image_i)
        if image_i.ndim == 2:
            image_i = image_i[..., np.newaxis]
        merged.append(image_i)
    merged = np.concatenate(merged, axis=-1)
    return merged