PythonでOpenCVとnumpyを使って画像の分割と結合

python コンピュータ
python

「SwinIR」というソフトで画像を拡大しようとしたところ、私の環境ではVRAMが不足で実行できませんでした。
同梱の小さめの画像は成功するので、画像を分割して「SwinIR」を実行し、その結果を再結合すれば目的を達成できるのではと思い「画像を分割と結合」するスクリプトを作成してみました。

分割スクリプト

import sys, os, glob
import cv2
import numpy as np

# 分割の幅と高さ
split_width = 256
split_height = 256
# ディレクトリ
in_dir = "./input"
out_dir = "./output"
# コマンドライン引数
args = sys.argv
if (len(args)>=3):
    in_dir = args[1]
    out_dir = args[2]
# 出力ディレクトリの作成
if not(os.path.exists(out_dir)):
    os.mkdir(out_dir)
# ファイルの一覧
for f in glob.glob((in_dir+'/*.*')):
    base, ext = os.path.splitext(os.path.basename(f))
    im = cv2.imread(f)
    h, w = im.shape[:2]
    #print([h, w])
    w_n = w // split_width
    if ((w % split_width) > 0):
        w_n = w_n + 1
    h_n = h // split_height
    if ((h % split_height) > 0):
        h_n = h_n + 1
    #print([h_n, w_n])
    for y in range(h_n):
        for x in range(w_n):
            yy = y*split_height
            xx = x*split_width
            height = split_height
            if (yy+height) > h:
                height = h - yy
            width = split_width
            if (xx+width) > w:
                width = w - xx
            tmp = im[yy:yy+height, xx:xx+width]
            #print(tmp.shape[:2])
            out_file = os.path.join(out_dir, "{0}_split_{1:03}_{2:03}{3}".format(base, y, x, ext))
            print(out_file)
            cv2.imwrite(out_file, tmp)

結合スクリプト

import sys, os, glob
import cv2
import numpy as np
import re

# 分割の幅と高さ
split_width = 256
split_height = 256
# ディレクトリ
in_dir = "./output"
out_dir = "./result"
# コマンドライン引数
args = sys.argv
if (len(args)>=3):
    in_dir = args[1]
    out_dir = args[2]
# 出力ディレクトリの作成
if not(os.path.exists(out_dir)):
    os.mkdir(out_dir)
# ファイルの一覧
basename = ""
im_h = None
im_v = None
yy = -1
xx = -1
for f in glob.glob((in_dir+'/*.*')):
    base, ext = os.path.splitext(os.path.basename(f))
    print(f)
    r = re.findall(r'(.+?)_split_(\d+)_(\d+)', base)
    if (r):
        b = r[0][0]
        y = int(r[0][1])
        x = int(r[0][2])
        print([b, y, x])
        tmp = cv2.imread(f)
        if (yy != y and im_h is not None):
            if (im_v is None):
                im_v = im_h.copy()
            else:
                im_v = cv2.vconcat([im_v, im_h])
            print(im_v.shape)
            im_h = None
        if (basename == ""):
            basename = b+ext
        if (basename != (b+ext) and im_v is not None):
            out_file = os.path.join(out_dir, basename)
            cv2.imwrite(out_file, im_v)
            basename = (b+ext)
            im_n = None
            im_v = None
            yy = -1
            xx = -1
        if (im_h is None):
            im_h = tmp.copy()
        else:
            im_h = cv2.hconcat([im_h, tmp])
        yy = y
        xx = x
if (basename != "" and im_v is not None):
    out_file = os.path.join(out_dir, basename)
    if (im_v is not None and im_h is not None):
        im_v = cv2.vconcat([im_v, im_h])
    print(im_v.shape)
    cv2.imwrite(out_file, im_v)

分割はnumpyの配列のスライスを使い、結合はOpenCVのvconcat()とhconcat()を使っています。
結合する画像ファイルはファイル名の昇順で読み込まれることを想定しています。
調べて動作確認しながら作成したのでスクリプトを全体を通しでみると、もう少し何とかなったのではないかとも思いますが、とりあえず動作したので良しとします。

ちなみに分割した画像を「SwinIR」で拡大し結合しなおした所、結合部分の違和感もなく結合していました。

出来上がった画像は「Real-ESRGAN」と似たような感じで、全体として少し曖昧な感じがする「Real-CUGAN」と比べると詳細部分までくっきりした感じです。拡大処理やノイズ除去としての性能とは異なりますが、おのおの生成される絵に異なった雰囲気があって比べると中々面白いと思います。ただ、処理のためのマシンパワーの確保が難点です。

追記20240106
SwinIRのメモリ不足ですが実行環境はVRAMが12GでOSはUbuntuでした。Windows環境で実行したところVRAMが8Gのグラフィックボードでも処理が完了しました。多分メインメモリ(共有メモリ?)を使っていると思われます。ただ1枚の画像を処理する(1000×1600→4000×6400)時間が数分ぐらい要するのでとても遅いです。

コメント