🔬

超解像を試してみた

2021/09/22に公開

Python

超解像

srgan

tech

下記のサイトと著者が公開されているGitHubのコードを使って超解像を試してみたのでまとめます。

Google Colabで動かせるようにしています。

変更点

Google Colab 用の cv2.imshow を使っています。

低解像度画像をもう少しモザイクっぽくするために、下記サイトを参考に一度かなり32サイズに変換したのち、64サイズに拡大しています。

高解像度画像
低解像度画像

かなりいい感じでボカせました。

class DownSizePairImageFolder(ImageFolder):
    def __init__(self, root, transform=None, large_size=256, middle_size=64, small_size=32, **kwds):
        super().__init__(root, transform=transform, **kwds)
        self.large_resizer = transforms.Resize(large_size)
        self.small_resizer = transforms.Resize(small_size) # 追加
        self.middle_resizer = transforms.Resize(middle_size)
        
    def __getitem__(self, index):
        path, _ = self.imgs[index]
        img = self.loader(path)
        large_img = self.large_resizer(img)
        small_img = self.small_resizer(img) # 追加
        middle_img = self.middle_resizer(small_img)
        if self.transform is not None:
            large_img = self.transform(large_img)
            small_img = self.transform(small_img) # 追加
            middle_img = self.transform(middle_img)
        return middle_img, large_img

エラー回避のため、下記のような微妙な修正をしています。

G_loss+=G_loss.data[0] # 修正前。エラーになる

G_loss+=G_loss.data # 修正後。エラーを回避できる

Google Drive のフォルダ構成

下記のとおりです。

なお、trainフォルダ、testフォルダ直下に、image フォルダを作成して画像を保存しているのは、trainフォルダ、testフォルダ直下に保存するとDataLoderで読みこめないからです。

image --- celeba --- asset 学習した重みを保存するフォルダ  
                  |
                  -- img_align_celeba データセットを保存するフォルダ  
                  |
                  -- save_image 訓練実施中に画像を保存するフォルダ  
                  |
                  -- test --
                  |        |
                  |        --- image img_align_celeba から任意の数のテスト用のデータセットを保存するフォルダ
                  -- train --   
                            |
                            -- image img_align_celeba から任意の数の訓練用のデータセットを保存するフォルダ

なお、save_image、asset、train / test 及び直下のimage フォルダは下記のコードで自動作成するようになっています。

if not os.path.exists("save_image"):
    os.mkdir("save_image")
if not os.path.exists("asset"):
    os.mkdir("asset")
if not os.path.exists("train"):
    os.mkdir("train")
if not os.path.exists("train/image"):
    os.mkdir("train/image")
if not os.path.exists("test"):
    os.mkdir("test")
if not os.path.exists("test/image"):
    os.mkdir("test/image")

留意点

pytorch のバージョンが新しいとエラーが起きます。

pytorch 1.4.0、torchvision 0.5.0 に落としてようやく動きました。

CelebAのデータセット、img_align_celeba.zipをGoogle Driveに入れて使っています。

データが多すぎるせいか、Google Drive内の画像の読込を失敗することがありますが、再実行すると読み込めます。
(自分の場合は1回目は必ず読込失敗しました)

調整前

first_train のエポック数を100回、学習済みモデルをVGG16、trainのエポック数を10,000回にしました。

・・・うーん微妙。

比較するとこんな感じ。
(左が元の画像、真ん中が低解像度画像、右がGANで生成した画像)

調整後

first_train のエポック数を10,000回に増やし、学習済みモデルを mobilenet_v2に変更にしました。

なお、trainのエポック数は据え置きの10,000回のままです。

35枚の画像を学習させて、約3.5時間ほどかかりました。

ちょっと崩れているものの、高解像度っぽい画像が生成されました！

やったぜ。

比較するとこんな感じ。
(左が元の画像、真ん中が低解像度画像、右がGANで生成した画像)

以上になります、最後までお読みいただきありがとうございました。

変更点

Google Drive のフォルダ構成

留意点

調整前

調整後

Discussion