【Pytorch】rgb转lab颜色空间转换

最近有需要，自己按照opencv实现的方式，自测！

个别tensor，numpy互转参考：https://oldpan.me/archives/pytorch-tensor-image-transform

颜色空间转换代码：

from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import os
import cv2
from color_convertion import *
import os

def tensor_to_np(tensor):
    img = tensor.mul(255).byte()
    img = img.cpu().numpy().squeeze(0).transpose((1, 2, 0))
    return img

def toTensor(img):
    assert type(img) == np.ndarray,'the img type is {}, but ndarry expected'.format(type(img))
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = torch.from_numpy(img.transpose((2, 0, 1)))
    return img.float().div(255).unsqueeze(0)  

# N x H x W X C numpy to tensor
def N_imgtoTensor(img):
    img = torch.from_numpy(img.transpose((0, 3, 1, 2)))
    return img.float().div(255)

def N_tensor_to_np(tensor): #tensor value range is [0,1]
    img = tensor.mul(255).byte()         #tensor size [N,C,H,w]
    img = img.cpu().numpy().transpose((0,2,3,1))
    return img     #np_img size [N,H,w,C]


def read_file(sn,tn):
    s = cv2.imread('some_data/'+sn+'.jpg')
    s = cv2.cvtColor(s,cv2.COLOR_BGR2LAB)
    t = cv2.imread('some_data/'+tn+'.jpg')
    t = cv2.cvtColor(t,cv2.COLOR_BGR2LAB)
    return s, t

def tensor_rgbtolab(sn,tn):
    s = cv2.imread('some_data/'+sn+'.jpg')
    s = cv2.cvtColor(s,cv2.COLOR_BGR2LAB)
    t = cv2.imread('some_data/'+tn+'.jpg')
    t = cv2.cvtColor(t,cv2.COLOR_BGR2LAB)
    s= toTensor(s)
    t=toTensor(t)
    return tensor2im(s),tensor2im(t)

def tensor_original_img(sn,tn):
    s = cv2.imread('some_data/'+sn+'.jpg')
    t = cv2.imread('some_data/'+tn+'.jpg')
    s= toTensor(s)
    t=toTensor(t)
    return tensor_to_np(s),tensor_to_np(t)


#1、 single img process!!!
s1,t1=read_file('p1','p2')

s2,t2=tensor_rgbtolab('p1','p2')

s3,t3=tensor_original_img('p1','p2')


cv2.namedWindow("1",0)
cv2.namedWindow("2",0)
cv2.namedWindow("3",0)
cv2.resizeWindow("1",312,600)
cv2.resizeWindow("2",312,600)
cv2.resizeWindow("3",312,600)
cv2.imshow("1",np.vstack([s1,t1]))
cv2.imshow("2",np.vstack([s2,t2]))
cv2.imshow("3",np.vstack([s3,t3]))



#2、multi batch to process!!!   
def tensor_rgbtolab(sn,tn):
    s = cv2.imread('some_data/'+sn+'.jpg')
    t = cv2.imread('some_data/'+tn+'.jpg')
    temp=[]
    for i,img in enumerate([s,t]):
        img = cv2.cvtColor(img,cv2.COLOR_BGR2LAB)
        s=np.expand_dims(img,axis=0)
        temp.append(s)

    cat_result=np.concatenate(temp,axis=0)
    N_tensor=N_imgtoTensor(cat_result)
    print(N_tensor.shape)
    result=N_tensor_to_np(N_tensor)

    return result

oo=tensor_rgbtolab('p1','p2')
for i,img in enumerate(oo):
    cv2.imshow(str(i),img)




cv2.waitKey(0)
cv2.destroyAllWindows()