PyTorch로 이미지 채널(RGB)별 평균 구하는 방법

RGB 이미지는 빨강(Red), 초록(Green), 파랑(Blue)의 세 가지 색상 채널로 구성됩니다. 이미지 처리 과정에서는 채널별 픽셀 값의 평균을 계산해야 하는 경우가 자주 있습니다. 이때 torch.mean() 메서드를 사용하면 되는데, 이 메서드의 입력값은 반드시 PyTorch 텐서여야 합니다. 따라서 먼저 이미지를 PyTorch 텐서로 변환한 뒤 메서드를 적용해야 합니다. torch.mean()은 텐서 내 모든 요소의 평균값을 반환하며, 이미지 채널 방향으로 평균을 구하려면 매개변수를 dim = [1,2]로 설정하면 됩니다.

구현 단계

필요한 라이브러리를 임포트합니다. 아래 예제에서는 torch, torchvision, Pillow, OpenCV를 사용하므로 사전에 설치되어 있어야 합니다.
Image.open() 함수로 입력 이미지를 읽어 변수 "img"에 할당합니다.
PIL 이미지를 PyTorch 텐서로 변환하기 위한 transform을 정의합니다.
정의한 transform을 사용해 이미지를 PyTorch 텐서로 변환하고, 그 결과를 "imgTensor"에 할당합니다.
torch.mean(imgTensor, dim = [1,2])를 계산합니다. 이 연산은 세 개의 값을 가진 텐서를 반환하며, 각각의 값이 RGB 세 채널의 평균입니다. 반환된 값은 "R_mean", "G_mean", "B_mean" 변수에 각각 할당할 수 있습니다.
이미지 픽셀의 세 채널 평균값을 출력합니다.

입력 이미지

두 예제 모두 아래 이미지를 입력으로 사용합니다.

PyTorch로 이미지 채널(RGB)별 평균 구하는 방법

예제 1: PIL 이미지 활용

# Python program to find mean across the image channels
# import necessary libraries
import torch
from PIL import Image
import torchvision.transforms as transforms

# Read the input image
img = Image.open('opera.jpg')

# Define transform to convert the image to PyTorch Tensor
transform = transforms.ToTensor()

# Convert image to PyTorch Tensor (Image Tensor)
imgTensor = transform(img)
print("Shape of Image Tensor:\n", imgTensor.shape)

# Compute mean of the Image Tensor across image channels RGB
R_mean, G_mean ,B_mean = torch.mean(imgTensor, dim = [1,2])

# print mean across image channel RGB
print("Mean across Red channel:", R_mean)
print("Mean across Green channel:", G_mean)
print("Mean across Blue channel:", B_mean)

출력 결과

Shape of Image Tensor:
    torch.Size([3, 447, 640])
Mean across Red channel: tensor(0.1487)
Mean across Green channel: tensor(0.1607)
Mean across Blue channel: tensor(0.2521)

예제 2: OpenCV 활용

OpenCV를 사용해서도 이미지를 읽을 수 있습니다. OpenCV로 읽은 이미지는 numpy.ndarray 타입이라는 점에 유의하세요. 이번 예제에서는 평균을 계산하는 또 다른 방법을 소개합니다. 바로 텐서의 기본 연산인 imgTensor.mean()을 사용하는 방식입니다. 아래 예제를 확인해 보세요.

# Python program to find mean across the image channels
# import necessary libraries
import torch
import cv2
import torchvision.transforms as transforms

# Read the input image either using cv2 or PIL
img = cv2.imread('opera.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Define transform to convert the image to PyTorch Tensor
transform = transforms.ToTensor()

# Convert image to PyTorch Tensor (Image Tensor)
imgTensor = transform(img)
print("Shape of Image Tensor:\n", imgTensor.shape)

# compute mean of the Image Tensor across image channels RGB
# The other way to compute the mean
R_mean, G_mean ,B_mean = imgTensor.mean(dim = [1,2])

# print mean across image channel RGB
print("Mean across Red channel:", R_mean)
print("Mean across Green channel:", G_mean)
print("Mean across Blue channel:", B_mean)

출력 결과

Shape of Image Tensor:
    torch.Size([3, 447, 640])
Mean across Red channel: tensor(0.1487)
Mean across Green channel: tensor(0.1607)
Mean across Blue channel: tensor(0.2521)