Image Convert Function

0. Package

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torchvision.transforms.functional as F


1. CV2 Image -> PLT Image (Print)
(100,200,3 : BGR) -> (100,200,3 : RGB)

# imshow function (cv2 image -> plt image -> plt.show)
# cv2 image를 plt image로 출력 
def imshow(in_img):
    out_img = cv2.cvtColor(in_img,cv2.COLOR_BGR2RGB) 
    plt.figure(figsize=(5,5))
    plt.imshow(out_img)
    plt.axis('off')
    plt.show()


2. Tensor Image -> PLT Image (Print)
(3,64,64 : RGB) -> (64,64,3 : RGB)

def timshow(img):
    img=img.permute(1,2,0)
    plt.figure(figsize=(5,5))
    plt.imshow(img)
    plt.axis('off')
    plt.show()


3. Normalized Tensor -> CV2 Image (Convert)
(3,64,64 : RGB, Normalized), (mean,std : RGB) -> (64,64,3 : BGR)

# RGB로 구한 평균, 표준편차 -> BGR로 구한 평균, 표준편차
def cv2_infor(tmean,tstd):
    cmean,cstd=tmean.clone(),tstd.clone()
    cmean[0],cmean[2]=tmean[2],tmean[0]
    cstd[0],cstd[2]=tstd[2],tstd[0]

    return cmean,cstd

def tensor_to_cv_img(data, mean, std):
    # RGB mean, std -> BGR mean,std
    mean,std=cv2_infor(mean,std)
    MEAN, STD = mean.numpy(),std.numpy()

    # tensor to cv
    img = data # Tensor Shape. [C, H, W]
    img = img.detach().cpu().numpy() # tensor -> numpy
    img = np.transpose(img, (1, 2, 0)) # [C,H,W] -> [H,W,C]
    img = np.clip(255.0 * (img * STD + MEAN), 0, 255) # denomalize
    img = img.astype(np.uint8).copy() # np.float32 -> np.uint8

    return img


4. CV2 Image -> Normalized Tensor (Convert)
(100,200,3 : BGR), (mean=0.5, std=0.5 : RGB) -> (3,64,64 : RGB, Normalized)

def cv2img_to_tensor(src):
    # resizing
    img=cv2.resize(src, (64, 64), interpolation=cv2.INTER_AREA)
    # cv2 to plt
    img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    # plt to tensor and normalization
    torch_img=F.to_tensor(img)
    torch_img=F.normalize(torch_img, mean=0.5, std=0.5)
    # edit shape
    output=torch_img.view(3,64,64).to(device)

    return output


Example Code: [GitHub]

Leave a Reply

Your email address will not be published. Required fields are marked *