【Pytorch】将文件夹下的所有文件依次通过模型

描述：

在合成视频时，需要将所有图片从视频里抽帧出来，经过模型处理，再把所有帧合成为视频，以下代码主要为将文件夹下所有文件通过模型的处理代码。

---

具体代码：

#所用到的库
import torch
import cv2
import os
import imageio
#模型相关文件
from option import args
from model import Model

device = torch.device("cuda",0)

def pic2target(pic_path,target_path):
    model = arbrcan.Model(args)
    model.load_state_dict(torch.load('/data1/video/model_180.pt', map_location='cpu'), strict=True)
    model.to(device)
    model.eval()
    model.set_scale(scale, scale2)
    #得到文件夹下所有文件
    filelist = os.listdir(pic_path)

    for i in filelist:
        load_path = pic_path+str(i)
        lr = imageio.imread(load_path)
        lr = torch.tensor(lr)
        lr = lr.unsqueeze(dim=0)
        lr = lr.permute(0,3,1,2)
        lr = lr.type(torch.FloatTensor).to(device)
        with torch.no_grad():
            hr = model(lr)
            hr = hr.squeeze(dim=0).permute(1,2,0)
            #转换为RGB格式
            hr = torch.cat([hr[:,:,2][:,:,None],hr[:,:,1][:,:,None],hr[:,:,0][:,:,None]],dim=2).cpu().detach().numpy()
        cv2.imwrite(target_path+str(i), hr)
        print('当前图片：',i)

if __name__ == '__main__':
    pic2target(pic_path = "/data1/video/tran5/",target_path='/data1/video/target5')