pytorch dataset에서 이미지를 읽어오고 변환하고 저장하는 과정
3채널 이미지 VS. 1채널 이미지
3채널 이미지의 경우 open cv로 읽어오는데
open cv의 경우 pil과는 다르게
BGR로 이미지를 읽어오기 때문에 RGB로 변환하는 과정이 필요하다. (cv2.COLOR_BGR2RGB)
1채널 이미지의 경우 IMREAD_GRAYSCALE로 읽어오면 변환 끝
읽어온 후의 이미지 shape을 찍어보면,
3채널의 경우 height, width, channel 3 dim인 반면
1채널의 경우 height, width 만 존재하고 channel이 없는 걸로 표현된다
따라서 1채널 이미지의 경우 channel의 dim을 만들어 줘야하는데
여기서는 2가지 방법이 존재한다.
1. np.expand_dims 를 사용해서 dim 하나를 추가한다 (albumentation or torchvision transform 둘 중 하나를 나중에 적용)
2. albumentation augmentation 을 사용할 거라면 albumentations.pytorch.functional 에 있는 img_to_tensor 함수를 사용한다.
1. np.expand_dims 를 사용
np.expand_dims 를 위와 같이 적용하면 dim이 1로 하나 증가한 것을 알 수 있다.
이 상태가 되면 3채널과 3번째 dim의 개수만 차이가 나고 동일한 조건이 되었다.
그 다음에는 torch model이 받는 input shape을 고려해야 하는데
torch의 input shape : batch, channel, height, width
이렇게 되기 때문에, batch는 dataloader에 넣으면 자동으로 변환되기 때문에 고려하지 않아도 되고
위에서 표현된 h, w, c 이 모양을 c, h, w으로 바꿔 줘야한다.
이 때 np.transpose를 사용해서 axis 를 바꾸면 된다.
np.transpose를 다음과 같이 사용해서 2번째 축인 channel이 가장 처음에 오도록 하고 h, w의 순서는 그대로 둔 것이다.
이렇게 되면 torch model의 input으로 넣을 수 있는 조건이 충족되고 augmentation을 진행하면 된다.
2. albumentations.pytorch.functional 에 있는 img_to_tensor 함수 사용
단, 이 함수를 사용하게 되면 이미 tensor로 변환한 것이기 때문에
albumentations.transforms 를 사용한 augmentation에서 ToTensor와 Normalize를 사용하지 않아도 된다.
앞서 말한 것처럼 image를 읽고 3채널의 경우 BGR2RGB를 사용해서 RGB로 가져오고, 1채널의 경우 np.expand_dims를 사용해서 dim 을 맞춰 준다음
img_to_tensor 함수를 적용하면 위 처럼 torch tensor이면서 c, h, w 순으로 적절하게 shape이 만들어진 것을 확인할 수 있다.
댓글