SRCNN(超解像)の学習させて保存したモデルを用いて，複数画像に超解像を適応するプログラム

###挨拶
超解像初学者です．読みにくいコードかもしれませんが，よろしくお願いします．

###やりたいこと
SRCNN超解像で学習して保存したモデルを用いて，超解像を複数画像に適応させたい．適応する画像は学習画像に比べると大きいものを想定しています．モデル適応時に画像を学習画像と同じサイズにクロップしてリストに保存し，それを超解像した後に画像を結合して大きな画像を超解像しようとしています．これを画像枚数分繰り返します．

###現状
今回は，33×33の学習画像で学習をしたモデルを用いて，超解像適応プログラムを組んでみました．すると，下記のようなエラーが出てしましました．
調べても解決法がわかりませんでした．
ほかにも，間違っている点を指摘してくださるとありがたいです．

作成したプログラム

import torch
import torchvision as tv
import torch.nn as nn
import numpy as np
import os
import time
import cv2
from natsort import natsorted
from PIL import Image
from torch.nn.functional import relu
from torchsummary import summary



#モデル
class SRCNN(nn.Module):
    def __init__(self):
        super(SRCNN, self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=9, padding=4)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=1, padding=0)
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=3, kernel_size=5, padding=2)
        
     #モデル全体の記述
    def forward(self, INPUT):
        CONV1 = self.conv1(INPUT)
        RELU1 = relu(CONV1)
        CONV2 = self.conv2(RELU1)
        RELU2 = relu(CONV2)
        CONV3 = self.conv3(RELU2)
         
        return CONV3
    

        

def main():
    
    input_dir = r"C:\Users\81703\Super_resolution\test\cal\1"
    save_dir = r"C:\Users\81703\Super_resolution\test\cal_sr\1"
    model_dir = r"C:\Users\81703\Super_resolution\SRCNN\dataset\model"
    train_dir = r"C:\Users\81703\Super_resolution\SRCNN\dataset\train"
    
    #モデルの読み込みと適応画像のリスト取得等々
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SRCNN().to(device)
    model.load_state_dict(torch.load(os.path.join(model_dir,"model_SRCNN.pth"), map_location=device))
    sample_img = Image.open(os.path.join(train_dir, "0.png"))
    summary(model,(3,sample_img.height,sample_img.width))
    start = time.time()
    apply_imgs_list = os.listdir(input_dir)
    apply_imgs_list = natsorted(apply_imgs_list)
    
    model.eval()
    with torch.no_grad():
        for n in apply_imgs_list:
        
            #画像の読み込みと名前，拡張子の取得
            os.chdir(input_dir)
            apply_img = Image.open(n).convert("RGB")
            img_name, img_ext = os.path.splitext(n)
            print(img_name)
            
            #画像サイズとクロップ数の計算部分
            numX = apply_img.width // sample_img.width
            numY = apply_img.height // sample_img.height
            crop_imgs = []
            out_imgs = []
            
            #画像を分割
            for i in range(numY):
                for j in range(numX):
                    input_img = apply_img.crop((sample_img.width * j, sample_img.height*i,
                                               sample_img.width * j + sample_img.width, sample_img.height * i + sample_img.height))
                    input_img_tensor = tv.transforms.ToTensor()(input_img)
                    crop_imgs.append(input_img_tensor)
                    
            #分割をモデルに適用
            for m in crop_imgs:
                prediction = model(m)
                out_imgs.append(prediction)
                
            #モデルの出力画像を繋げる
            append_imgs = out_imgs
            v_img = []
            for y in range(numY):
                u_img = []
                for x in range(numX):
                    num = x + y * numX
                    u_img.append(append_imgs[num])
                
                imgU = cv2.hconcat(u_img)
                v_img.append(imgU)
                
            append_img = cv2.vconcat(v_img)
            append_img = Image.fromarray(np.unit8(append_img))
            save_name = str(n)
            save_name_dir = os.path.join(save_dir, save_name)
            append_img.save(save_name_dir)
    end = time.time() - start
    print(end)
    print("finish")

if __name__ == "__main__":
    main()

エラー内容

RuntimeError: Expected 4-dimensional input for 4-dimensional weight 64 3 9, but got 3-dimensional input of size [3, 33, 33] instead

行動規範の内容に同意します

回答1件

エラーメッセージは入力が4次元を想定しているのに、3次元の入力が入ってきた
という意味です。

pytorchの場合、モデルに入力するテンソルの次元は(Batch_Size, Channel, Height, Width)の4次元です。
(https://discuss.pytorch.org/t/dimensions-of-an-input-image/19439)
ですが、ご提示いただいているコードでは入力が(Channel, Height, Width)の3次元となっているため、エラーが起きていると考えられます。
つまり、入力のtensorに新しい次元を足せばよいということになります。

また、画像の訓練時のサイズとデプロイ時のサイズは必ずしも一致している必要なないかと思います。
SRGANなどの学習でも,(64x64)のパッチで学習を行います。
推論時では、元の画像の大きさをそのまま入力するやり方を取っていたはずです。

そのため、ご提示いただいていたコードであるような、

python
1            #画像の読み込みと名前，拡張子の取得
2            os.chdir(input_dir)
3            apply_img = Image.open(n).convert("RGB")
4            img_name, img_ext = os.path.splitext(n)
5            print(img_name)
6
7            #画像サイズとクロップ数の計算部分
8            numX = apply_img.width // sample_img.width
9            numY = apply_img.height // sample_img.height
10            crop_imgs = []
11            out_imgs = []
12
13            #画像を分割
14            for i in range(numY):
15                for j in range(numX):
16                    input_img = apply_img.crop((sample_img.width * j, sample_img.height*i,
17                                               sample_img.width * j + sample_img.width, sample_img.height * i + sample_img.height))
18                    input_img_tensor = tv.transforms.ToTensor()(input_img)
19                    crop_imgs.append(input_img_tensor)
20
21            #分割をモデルに適用
22            for m in crop_imgs:
23                prediction = model(m)
24                out_imgs.append(prediction)
25
26            #モデルの出力画像を繋げる
27            append_imgs = out_imgs
28            v_img = []
29            for y in range(numY):
30                u_img = []
31                for x in range(numX):
32                    num = x + y * numX
33                    u_img.append(append_imgs[num])
34
35                imgU = cv2.hconcat(u_img)
36                v_img.append(imgU)
37
38            append_img = cv2.vconcat(v_img)
39            append_img = Image.fromarray(np.unit8(append_img))
40            save_name = str(n)
41            save_name_dir = os.path.join(save_dir, save_name)
42            append_img.save(save_name_dir)

のような部分は

apply_img = Image.open(n).convert("RGB")
input_image_tensor = tv.transforms.ToTensor()(apply_img)
input_image_tensor = torch.unsqueeze(input_image_tensor, 0)
prediction = model(m)

としても良いアウトプットを得られるかと思います。
SRGANなどの学習でも,(64x64)のパッチで学習を行い、デプロイ時はパッチを利用しないで推論していたと思います。

投稿2021/04/09 05:05

編集2021/04/09 05:19