pytorch图像处理的问题

img = cv2.imread(os.path.join(path, file, pic))
img = Image.fromarray(img.astype(np.uint8))

img = Image.open(os.path.join(path, file, pic))

#coding:utf-8
import torch.nn as nn
import torch

def dwt_init(x):

    x01 = x[:, :, 0::2, :] / 2
    x02 = x[:, :, 1::2, :] / 2
    x1 = x01[:, :, :, 0::2]
    x2 = x02[:, :, :, 0::2]
    x3 = x01[:, :, :, 1::2]
    x4 = x02[:, :, :, 1::2]
    x_LL = x1 + x2 + x3 + x4
    x_HL = -x1 - x2 + x3 + x4
    x_LH = -x1 + x2 - x3 + x4
    x_HH = x1 - x2 - x3 + x4

    return torch.cat((x_LL, x_HL, x_LH, x_HH), 1)


# 使用哈尔 haar 小波变换来实现二维逆向离散小波
def iwt_init(x):
    r = 2
    in_batch, in_channel, in_height, in_width = x.size()
    #print([in_batch, in_channel, in_height, in_width])
    out_batch, out_channel, out_height, out_width = in_batch, int(
        in_channel / (r**2)), r * in_height, r * in_width
    print(out_batch, out_channel, out_height, out_width)
    x1 = x[:, 0:out_channel, :, :] / 2
    x2 = x[:, out_channel:out_channel * 2, :, :] / 2
    x3 = x[:, out_channel * 2:out_channel * 3, :, :] / 2
    x4 = x[:, out_channel * 3:out_channel * 4, :, :] / 2
    print(x1.shape)
    print(x2.shape)
    print(x3.shape)
    print(x4.shape)
    # h = torch.zeros([out_batch, out_channel, out_height, out_width]).float().cuda()
    h = torch.zeros([out_batch, out_channel, out_height, out_width]).float()

    h[:, :, 0::2, 0::2] = x1 - x2 - x3 + x4
    h[:, :, 1::2, 0::2] = x1 - x2 + x3 - x4
    h[:, :, 0::2, 1::2] = x1 + x2 - x3 - x4
    h[:, :, 1::2, 1::2] = x1 + x2 + x3 + x4

    return h


# 二维离散小波
class DWT(nn.Module):
    def __init__(self):
        super(DWT, self).__init__()
        self.requires_grad = False  # 信号处理，非卷积运算，不需要进行梯度求导

    def forward(self, x):
        return dwt_init(x)


# 逆向二维离散小波
class IWT(nn.Module):
    def __init__(self):
        super(IWT, self).__init__()
        self.requires_grad = False

    def forward(self, x):
        return iwt_init(x)

if __name__ == \'__main__\':
    import os, cv2, torchvision
    from PIL import Image
    import numpy as np
    from torchvision import transforms as trans
    # img = cv2.imread(\'./1.jpg\')
    # print(img.shape)
    # img = Image.fromarray(img.astype(np.uint8))
    img = Image.open(\'./1.jpg\')
    transform = trans.Compose([
        trans.ToTensor()
    ])
    img = transform(img).unsqueeze(0)
    dwt = DWT()
    change_img_tensor = dwt(img)
    print(change_img_tensor.shape)

    for i in range(change_img_tensor.size(1)//3):
        print(i*3,i*3+3)
        torchvision.utils.save_image(change_img_tensor[:,i*3:i*3+3:,:], os.path.join(\'./\', \'change_{}.jpg\'.format(i)))

from pathlib import Path
from config import get_config
import argparse

from PIL import Image
from tqdm import tqdm
import mxnet as mx
import cv2
import numpy as np

def load_mx_rec(rec_path):
    save_path = rec_path/\'imgs\'
    if not save_path.exists():
        save_path.mkdir()
    imgrec = mx.recordio.MXIndexedRecordIO(str(rec_path/\'train.idx\'), str(rec_path/\'train.rec\'), \'r\')
    img_info = imgrec.read_idx(0)
    header,_ = mx.recordio.unpack(img_info)
    max_idx = int(header.label[0])
    for idx in tqdm(range(1,max_idx)):
        img_info = imgrec.read_idx(idx)
        header, img = mx.recordio.unpack_img(img_info)
        label = int(header.label)
        #add
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        img = Image.fromarray(img)
        label_path = save_path/str(label)
        if not label_path.exists():
            label_path.mkdir()
        img.save(label_path/\'{}.jpg\'.format(idx), quality=95)

if __name__ == \'__main__\':
    parser = argparse.ArgumentParser(description=\'for face verification\')
    parser.add_argument("-r", "--rec_path", help="mxnet record file path",default=\'faces_emore\', type=str)
    args = parser.parse_args()
    conf = get_config()
    rec_path = conf.data_path/args.rec_path
    load_mx_rec(rec_path) #训练数据在train.idx和train.rec

img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)