build_notta_tta_comp_db.py 初心者向け解説

前のページでは「動かし方」を説明しました。このページは、このコードの仕組みそのもの——なぜこう書かれているのか——を理解したい方向けです。このプロジェクトの前提知識は不要です。 DOI 1.19.1 として無編集のまま全文公開されています。

各ライブラリが何をしているか

このスクリプトは6つの外部ライブラリに依存しています（公開物の requirements.txt参照）。名前だけでなく、このコードの中で実際に何をしているかを説明します：

`numpy`	大量の数値をまとめて効率よく扱うための基本ライブラリ。画像は突き詰めると「画素の明るさの数値がずらっと並んだもの」なので、numpyを使うと8×8ブロック(あるいは画像全体)を一気に計算でき、1画素ずつループで処理するより格段に速い。
`opencv-python-headless`（`cv2`としてimport）	広く使われているコンピュータビジョン用ライブラリ。「headless」は画面表示機能を持たないという意味——このスクリプトは画像を画面に表示する必要が一切ないので、ちょうど良い選択。（`heatmap_tool.py`内で）鮮明さ・ノイズレベルといった特定の計算に使用。
`scikit-image`（`skimage`の一部としてimport）	もう1つの画像処理ライブラリ。ここでは特定のテクスチャ解析手法（「グレーレベル同時生起行列」——隣り合う画素の明るさの組み合わせがどれだけ頻繁に現れるかを見る、テクスチャの「複雑さ」を数値化する標準的な方法）専用に使用。
`pillow`（`PIL`としてimport）	PNGやJPEGなどの画像ファイルを開いて、他のライブラリが扱える画素データに変換するための標準ライブラリ。
`duckdb`	別途サーバーを立てる必要のない、ローカルで動くデータベースエンジン。計算した全ての数値が最終的に格納される場所であり、後からSQLで直接調査できる。
`pandas`	表形式のデータ（行と列、スプレッドシートのようなもの）をメモリ上で扱うライブラリ。ここでは純粋に中継ぎ役として使う——1画像分の数万個の数値を計算したら、まず表としてまとめ、それをduckdbへ一括で書き込む——1行ずつ書き込むよりずっと速い。

コードをセクションごとに見る

どのファイル同士が対応するかを見つける

1つの画像は3つの別々のファイル——ground truth版、TTAなしで復元した版（notta）、TTAありで復元した版（tta）——として存在し、それぞれ別のフォルダに、しかもモデルごとに微妙に違う命名規則で入っていることが多いです。このスクリプトの最初の仕事は、画像IDごとに、3つのファイルが実際に全て揃っているかを調べることです。

def build_lookup(model_dir):
    lookup = {}
    for path in glob.glob(os.path.join(model_dir, '*.png')):
        base = os.path.basename(path)
        m = IMAGE_ID_RE.match(base)
        if m:
            image_id = m.group(1)
        elif '_LR' in base:
            image_id = base.split('_LR')[0]
        else:
            image_id = os.path.splitext(base)[0]
        lookup[image_id] = path
    return lookup

この関数はフォルダの中身を1回だけリストアップし、（画像ID→ファイルパスの）対応表をメモリ上に作ります。逆に「画像1枚ごとに毎回フォルダを検索し直す」というやり方は一見問題なさそうに見えますが、フォルダに数千個のファイルがあると、これが数千回分のフォルダ全体スキャンになり、画像が増えるほど劇的に遅くなります——特にネットワーク越しやマウントされたドライブでは顕著です。1回だけリストアップして、以降はメモリ上の対応表を引くだけにすることで、この問題を完全に回避しています。

画像1枚あたりのメインの計算処理

これが、画像1枚分の実際の計算を行う関数です。全ての画像に対して独立に、並列に（詳細は後述）実行されます。

def process_one_image(args):
    image_id, gt_path, notta_path, tta_path = args
    try:
        gt_u8, gt_f = load_gray(gt_path)
        notta_u8, notta_f = load_gray(notta_path)
        tta_u8, tta_f = load_gray(tta_path)
        if not (gt_u8.shape == notta_u8.shape == tta_u8.shape):
            return image_id, None, f'size mismatch ...'
        h, w = gt_u8.shape
        if h % BLOCK_SIZE or w % BLOCK_SIZE:
            return image_id, None, f'size {h}x{w} not a multiple of {BLOCK_SIZE}'

        n, grid_notta = compute_ssim_grid(gt_f, notta_f)
        _, grid_tta = compute_ssim_grid(gt_f, tta_f)
        comp_gt = compute_comp_grid(gt_u8, n)
        comp_notta = compute_comp_grid(notta_u8, n)
        comp_tta = compute_comp_grid(tta_u8, n)
        ...
        df = pd.DataFrame(data)
        img_info = (image_id, gt_path, notta_path, tta_path, w, h)
        return image_id, (img_info, df), None
    except Exception as e:
        return image_id, None, f'{type(e).__name__}: {e}'

順を追うと：3枚の画像を全てグレースケールの画素配列として読み込む（comp_*特徴量計算用に整数型、SSIM計算用に小数型の両方——load_grayが両方を返す）。3枚が実際に同じサイズであること、そのサイズが8×8ブロックで割り切れて余りが出ないことを確認する。その後、SSIMを2通り（ground truth 対notta、ground truth対tta）、8つのcomp_*特徴量を3通り（ground truth・notta・ttaそれぞれの画素から独立に）計算する ——これら全てが、画像内の全ての8×8ブロックに対して一気に行われます。1ブロックずつのループではなく、numpyの配列演算で書かれているためです。

関数全体がtry/exceptで囲まれている点にも注目してください：1枚の画像で何か問題が起きても（ファイルの破損、想定外のサイズなど）、その画像だけをエラーメッセージ付きでスキップし、数千枚の処理の途中で全体がクラッシュすることはありません。

データベースの構造を用意する

def ensure_schema(con):
    con.execute("""
        CREATE TABLE IF NOT EXISTS images (
            image_id VARCHAR PRIMARY KEY,
            path_gt VARCHAR, path_notta VARCHAR, path_tta VARCHAR,
            width INTEGER, height INTEGER
        );
    """)
    ...
    con.execute(f"""
        CREATE TABLE IF NOT EXISTS blocks (
            image_id VARCHAR, row INTEGER, col INTEGER,
            sigma_gt DOUBLE, mu_gt DOUBLE,
            ...
            PRIMARY KEY (image_id, row, col)
        );
    """)

これは、前のページで説明した2つのデータベーステーブルを、まだ存在しなければ作成します。ここで重要なのはIF NOT EXISTSです——これがあるおかげで、スクリプトを再実行するたびにこの関数を呼んでも安全であり、以前の（中断された）実行のデータを消してしまうことがありません。

全体をつなぐ：再開可能性と並列処理

前のページで説明した「再開できる」という仕組みが実際に実装されているのが、このmain()関数です：

done = set(r[0] for r in con.execute('SELECT image_id FROM images').fetchall())
...
for image_id in all_ids:
    if image_id in done:
        continue
    ...
    tasks.append((image_id, ...))

if args.limit is not None:
    tasks = tasks[:args.limit]

何か処理を始める前に、スクリプトはデータベースに既に記録済みの画像IDを読み込み、doneという集合に入れます。その直後のループは、この集合に既に含まれる画像を単純にスキップするだけです。再開可能性の仕組みはこれが全てで、「既にあるものを確認し、やり直さない」という以上の複雑さはありません。 --limitはこのフィルタリングの後に適用されるため、小さい--limitで同じコマンドを繰り返し実行すると、毎回新しいバッチの画像が処理され、同じ画像を再処理することはありません。

for chunk_start in range(0, len(tasks), args.chunk_size):
    chunk = tasks[chunk_start:chunk_start + args.chunk_size]
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        results = list(executor.map(process_one_image, chunk))

    for image_id, payload, error in results:
        ...
        con.execute('INSERT INTO images VALUES (?, ?, ?, ?, ?, ?)', img_info)
        con.append('blocks', df)
        n_done += 1

画像は小さなバッチ（「チャンク」、デフォルト20枚）単位で処理されます。1つのチャンクの中では、ProcessPoolExecutor が画像を複数のワーカープロセスに同時に振り分けます——これが、 1枚ずつ順番に処理するのではなく、複数のCPUコアを使えるようにしている部分です。1つのチャンク分の画像が全て処理し終わった時点で初めて、その結果をデータベースに書き込み、進捗行を表示します。これが、チャンクの途中で実行を中断するとそのチャンク分の進捗が失われる（何も書き込まれていないため）一方、既に完了したチャンクは安全に記録されたままである理由でもあります ——まさに実演動画で実際に起きている通りです。