[torchvision] torchvision.utils.save_image and torchvision.io.encode_jpeg, torchvision.io.encode_png

주로 다음의 3가지를 개인적으로 사용함.

• save_image: 파일 직접 저장, 배치 그리드 배열, 정규화/시각화 기능 내장
• encode_jpeg: 손실 압축, 더 작은 파일 크기, RGB/Grayscale만 지원
• encode_png: 무손실 압축, 더 큰 파일 크기, RGBA 지원 가능

관련 gist

https://gist.github.com/dsaint31x/c85521ce40c10787f04a7e5909373d22

dl_torchvision_save_images.ipynb

 

https://dsaint31.tistory.com/402

[DIP] Image Format (summary)

---

1. torchvision.utils.save_image

torchvision.utils 모듈의 save_image 를 통해 tensor 객체를 image 로 저장할 수 있음.

• 실은, 여러 이미지의 배치(batch)를 하나의 그리드(grid) 형태로 저장하는 데 주로 사용되는 함수임.
• 때문에 단일 이미지인 경우, nrow는 무시되나 padding은 여전히 영향을 줌.

torchvision.utils.save_image(
    tensor: torch.Tensor,
    fp: str | Path | BinaryIO,
    format: Optional[str] = None,
    *,
    nrow: int = 8,
    padding: int = 2,
    pad_value: float = 0.0,
    normalize: bool = False,
    value_range: Optional[Tuple[float, float]] = None,
    scale_each: bool = False,
)

• tensor :
  • 저장할 이미지 텐서 (dtype인 float32여야 함).
  • [C, H, W] or [N, C, H, W]
  • [0.0, 1.0] 이어야 함. (아니면 clamp 발생.)
• fp : 저장할 파일 경로 또는 파일 객체
• format : fp의 extension으로 자동으로 설정되나 지정도 가능. "png", "jpeg"
• nrow : batch 형태일 경우, 한 줄에 표시될 이미지 수.
• padding : 이미지 사이의 패딩 pixel 수.
• pad_value : padding으로 채울 값.
• normalize : value_rage 와 같이 정규화할 범위지정을 통한 정규화 수행여부.
• scale_each : batch의 각 이미지별로 개별 정규화.

[0.0,1.0] 사이의 float값으로 구성된 input을 가정하며, 해당 범위로 clamp 처리됨.

그리고 0.0과 1.0을 0과 255로 매핑하여 scaling함.

 

예제 코드

import torch
import torchvision.utils as utils
import torchvision.transforms as transforms

# 주석 처리를 해보고 돌려볼 것.
original_img = original_img.to(dtype=torch.float32)
original_img /= 255.

# 1. 단일 이미지 저장
utils.save_image(original_img.to(dtype=torch.float32), 'single.png')

# 2. 배치 이미지를 그리드로 저장
batch_images = torch.stack([original_img] * 16, dim=0)
print(f"{batch_images.shape = }")
# batch_images = torch.rand(16, 3, 64, 64)  # [N, C, H, W]
utils.save_image(
    batch_images.to(dtype=torch.float32), 
    'grid.png',
    nrow=4,          # 4x4 그리드
    padding=2,       # 이미지간 2픽셀 간격
    pad_value=1.0    # 흰색 패딩
)

# 3. 정규화가 필요한 경우 (-1~1 범위를 0~1로 처리하고 이를 다시 0-255로 저장)
normalized_images = torch.rand(8, 3, 32, 32) * 2 - 1  # [-1, 1] 범위
utils.save_image(
    normalized_images.to(dtype=torch.float32),
    'w_normalized.png',
    nrow=4,
    normalize=True,
    value_range=(-1, 1)    # 입력 범위 지정
)
utils.save_image(
    normalized_images.to(dtype=torch.float32),
    'wo_normalized.png',
    nrow=4,
    normalize=False,
    # range=(-1, 1)    # 입력 범위 지정
)

# 4. 각 이미지별 개별 정규화
varied_range_images = torch.stack([
    torch.rand(3, 32, 32) * 0.5,      # [0, 0.5] 범위
    torch.rand(3, 32, 32) * 2,        # [0, 2.0] 범위
    torch.rand(3, 32, 32) * 0.1 + 0.9 # [0.9, 1.0] 범위
])
print(f"{varied_range_images.shape = }")

utils.save_image(
    varied_range_images.to(dtype=torch.float32),
    'scale_each.png',
    nrow=3,
    scale_each=True  # 각각 개별 정규화
)

 

---

2. torchvision.io.encode_jpeg

torchvision.io 모듈의 함수로, tensor 객체를 jpeg으로 인코딩한 bytes객체를 반환함.

• 메모리에서 이미지 텐서를 JPEG 형식의 바이트 데이터로 변환하는 함수임.
• 파일로 저장하거나 네트워크 전송 등에 활용 가능함.

torchvision.io.encode_jpeg(
    input: torch.Tensor,
    quality: int = 75
) -> torch.Tensor

• input : 인코딩할 이미지 텐서.
  • [3, H, W] 또는 [1, H, W] 형태의 uint8 타입
  • 반드시 [0,255] 범위의 값이어야 함
• quality : JPEG 압축 품질. 1(최저품질)~100(최고품질) 범위

반환값: JPEG 바이트 데이터가 담긴 1차원 uint8 텐서

import torch
import torchvision.io as io

# uint8 텐서 준비 (0-255 범위)
tensor = torch.randint(0, 256, (3, 224, 224), dtype=torch.uint8)

# JPEG 바이트로 인코딩
jpeg_bytes = io.encode_jpeg(tensor, quality=90)

# 파일로 저장
with open('output.jpg', 'wb') as f:
    f.write(jpeg_bytes.numpy().tobytes())

torchvision.io.encode_png

torchvision.io 모듈의 encode_png를 통해 tensor 객체를 PNG 바이트(bytes객체)로 인코딩.

• 메모리에서 이미지 텐서를 PNG 형식의 바이트 데이터로 변환하는 함수임.
• 무손실 압축으로 정확한 픽셀 값 보존이 필요한 경우에 사용함.

torchvision.io.encode_png(
    input: torch.Tensor,
    compression_level: int = 6
) -> torch.Tensor

• input : 인코딩할 이미지 텐서. [3, H, W], [1, H, W], 또는 [4, H, W] 형태의 uint8 타입
• compression_level : PNG 압축 레벨. 0(압축없음)~9(최대압축) 범위

반환값: PNG 바이트 데이터가 담긴 1차원 uint8 텐서

import torch
import torchvision.io as io

# uint8 텐서 준비 (0-255 범위)
tensor = torch.randint(0, 256, (3, 224, 224), dtype=torch.uint8)

# PNG 바이트로 인코딩
png_bytes = io.encode_png(tensor, compression_level=6)

# 파일로 저장
with open('output.png', 'wb') as f:
    f.write(png_bytes.numpy().tobytes())

---

같이보면 좋은 자료

2025.06.17 - [Python] - [torchvision] image로부터 torch.tensor 객체 얻기

[torchvision] image로부터 torch.tensor 객체 얻기

2024.12.03 - [Python] - [Py] io.StringIO 와 io.BytesIO

[Py] io.StringIO 와 io.BytesIO

---