formatting.py

Go to the documentation of this file.

from typing import Sequence
 
import torch
import numpy as np
 
def to_tensor(data):
    """
    데이터를 텐서로 변환하는 기능
    args:
        data : 다양한 타입의 데이터
    return: Tensor로 변환된 데이터
    """
    if isinstance(data, torch.Tensor):
        return data
    elif isinstance(data, np.ndarray):
        return torch.from_numpy(data)
    elif isinstance(data, Sequence) and not isinstance(data, str):
        return torch.tensor(data)
    elif isinstance(data, int):
        return torch.LongTensor([data])
    elif isinstance(data, float):
        return torch.FloatTensor([data])
    elif isinstance(data, dict):
        for key in data.keys():
            data[key] = to_tensor(data[key])
    else:
        raise TypeError(f'to_tensor : type {type(data)} cannot be converted to tensor.')
 
class ToTensor(object):
    """
    딕셔너리 타입의 데이터에서 원하는 키의 데이터만 텐서로 변환하기 위한 클래스
    args:
        keys (list[str]) : 텐서 변환을 원하는 키, None이면 모든 키에 대해서 Tensor 변환
    """
    def __init__(self, keys=None):
        self.keys = keys
 
    def __call__(self, sample):
 
        if self.keys is None:
            sample = to_tensor(sample)
        else:
            for key in keys:
                assert key in sample, \
                    f'ToTensor : not found key in sample : {key}'
 
            for key in self.keys:
                samples[key] = to_tensor(samples[key])
 
        return sample
        
class ImageToTensor(object):
    """
    이미지를 텐서로 변환하기 위한 클래스
    args:
        keys (list[str]): 입력 데이터(dict)에서 이미지에 해당하는 키
    """
    def __init__(self, keys):
        self.keys = keys
 
    def transform(self, sample) -> dict:
        """
        입력 데이터의 이미지를 텐서로 변환하는 기능
        args:
            sample (dict): 데이터
        return: 텐서로 변환된 이미지가 포함된 데이터
        """
        #예외처리, 입력데이터에 설정된 키가 없으면 종료
        for key in keys:
            assert key in sample, \
                f'ImageToTensor : not found key in sample : {key}'
        
        #텐서 변환
        for key in self.keys:
            img = sample[key] #이미지
            if len(img.shape) < 3: #이미지가 2차원일 경우(흑백)
                img = np.expand_dims(img, -1) #끝에 한차원 늘려줘서 차원을 맞춰줌, (H, W, C(1))
            
            #텐서 변환 및 인덱스 구조 변경, H, W, C -> C, H, W
            sample[key] = (ToTensor.to_tensor(img.transpose(2, 0, 1))).contiguous()
 
        return sample
 
class CollectKeys(object):
    """
    데이터에서 원하는 데이터만 수집하기위한 클래스
    args:
        collect_keys (list[str]): 수집을 원하는 데이터의 키
    """
    def __init__(self, collect_keys):
        self.collect_keys = collect_keys
 
    def __call__(self, sample):
 
        #예외처리, 입력데이터에 설정된 키가 없으면 종료
        for key in collect_keys:
            assert key in sample, \
                f'CollectKeys : not found key in sample : {key}'
 
        #데이터 수집
        pack_data = {}
        for key in collect_keys:
            pack_data[key] = sample[key]
 
        return pack_data