深度画像からのステレオペア画像生成

# ステレオペア画像生成プログラム
#   カラー画像と深度画像からステレオビジョン用画像を生成
#   参考論文: Marr, D., & Poggio, T. (1976). Cooperative computation of stereo disparity. Science, 194(4262), 283-287.
#   GitHub: https://github.com/opencv/opencv
#   特徴: 深度マップから視差計算による立体視画像生成、三角測量の原理を応用
#         オクルージョン処理と穴埋めアルゴリズム

import cv2
import numpy as np
import tkinter as tk
from tkinter import filedialog

# カメラパラメータ定数
FOCAL_LENGTH = 800.0
BASELINE = 65.0
DEPTH_MIN = 500.0
DEPTH_MAX = 5000.0

def generate_confidence_map(depth_image, disparity, max_disparity):
    """深度値の信頼性に基づく信頼度マップを生成"""
    confidence = np.ones_like(depth_image, dtype=np.float32)

    # 深度勾配の計算（エッジ領域の検出）
    depth_float = depth_image.astype(np.float32)
    grad_x = cv2.Sobel(depth_float, cv2.CV_32F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(depth_float, cv2.CV_32F, 0, 1, ksize=3)
    gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)

    # エッジ領域は信頼度を下げる
    confidence = np.exp(-gradient_magnitude / 50.0)

    # 視差の妥当性チェック
    disparity_confidence = np.clip(1.0 - disparity / max_disparity, 0.1, 1.0)
    confidence *= disparity_confidence

    # 信頼度を0.1-1.0の範囲にクリップ
    confidence = np.clip(confidence, 0.1, 1.0)

    return confidence

def edge_preserving_inpainting(image, mask, depth_image):
    """エッジ保持を考慮したインペインティング"""
    if len(image.shape) == 3:
        height, width, channels = image.shape
        result = np.zeros_like(image)

        # チャンネルごとにインペインティング実行
        for c in range(channels):
            channel = image[:, :, c]
            result[:, :, c] = cv2.inpaint(channel, mask, 5, cv2.INPAINT_NS)
    else:
        result = cv2.inpaint(image, mask, 5, cv2.INPAINT_NS)

    # エッジ検出による構造保持
    edges = cv2.Canny(depth_image, 50, 150)

    # エッジ近傍では元の画像を優先的に使用
    edge_dilated = cv2.dilate(edges, np.ones((3,3), np.uint8), iterations=1)
    edge_mask = (edge_dilated > 0) & (mask == 0)

    if len(image.shape) == 3:
        edge_mask_3d = np.stack([edge_mask] * 3, axis=2)
        result = np.where(edge_mask_3d, image, result)
    else:
        result = np.where(edge_mask, image, result)

    return result

def uncertainty_based_blur(image, confidence_map):
    """信頼度に基づく適応的ガウシアンブラー"""
    # 信頼度が低い領域にブラーを適用
    blur_strength = (1.0 - confidence_map) * 3.0 + 1.0  # 1.0-4.0の範囲

    # 異なる強度のブラー画像を準備
    blurred_light = cv2.GaussianBlur(image, (5, 5), 1.0)
    blurred_medium = cv2.GaussianBlur(image, (7, 7), 2.0)
    blurred_heavy = cv2.GaussianBlur(image, (9, 9), 3.0)

    # 信頼度に基づく重み付け合成
    if len(image.shape) == 3:
        confidence_3d = np.stack([confidence_map] * 3, axis=2)
        high_conf_mask = confidence_3d > 0.8
        med_conf_mask = (confidence_3d > 0.5) & (confidence_3d <= 0.8)
        low_conf_mask = confidence_3d <= 0.5

        result = np.where(high_conf_mask, image,
                 np.where(med_conf_mask, blurred_light,
                 np.where(low_conf_mask, blurred_medium, blurred_heavy)))
    else:
        high_conf_mask = confidence_map > 0.8
        med_conf_mask = (confidence_map > 0.5) & (confidence_map <= 0.8)
        low_conf_mask = confidence_map <= 0.5

        result = np.where(high_conf_mask, image,
                 np.where(med_conf_mask, blurred_light,
                 np.where(low_conf_mask, blurred_medium, blurred_heavy)))

    return result.astype(np.uint8)

def improved_hole_filling(right_image, depth_image, disparity, max_disparity):
    """穴埋め処理"""

    # 1. 信頼度マップ生成
    confidence = generate_confidence_map(depth_image, disparity, max_disparity)

    # 2. 穴領域の検出
    if len(right_image.shape) == 3:
        mask = np.all(right_image == 0, axis=2).astype(np.uint8)
    else:
        mask = (right_image == 0).astype(np.uint8)

    # 穴がない場合は元の画像を返す
    if np.sum(mask) == 0:
        return right_image

    # 3. マスク領域の軽微な拡張（境界を滑らかに）
    kernel = np.ones((3,3), np.uint8)
    expanded_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # 4. エッジ保持インペインティング
    inpainted = edge_preserving_inpainting(right_image, expanded_mask, depth_image)

    # 5. バイラテラルフィルタによるエッジ保持平滑化
    if len(inpainted.shape) == 3:
        smoothed = cv2.bilateralFilter(inpainted, 9, 75, 75)
    else:
        smoothed = cv2.bilateralFilter(inpainted, 9, 75, 75)

    # 6. 信頼度に基づく適応的ブラー
    final_result = uncertainty_based_blur(smoothed, confidence)

    # 7. 元画像の有効部分を保持
    if len(right_image.shape) == 3:
        valid_mask = ~np.all(right_image == 0, axis=2)
        valid_mask_3d = np.stack([valid_mask] * 3, axis=2)
        final_result = np.where(valid_mask_3d, right_image, final_result)
    else:
        valid_mask = (right_image != 0)
        final_result = np.where(valid_mask, right_image, final_result)

    return final_result

# プログラム開始
print("========================================")
print("ステレオペア画像生成プログラム")
print("========================================")
print("\n【概要】")
print("カラー画像と深度画像から立体視用のステレオペア画像を生成します")
print("\n【操作方法】")
print("1. 深度表現方式を選択（A または B を入力）")
print("2. カラー画像ファイルを選択")
print("3. 深度画像ファイルを選択")
print("4. 処理完了後、画像が表示されます")
print("5. 任意のキーを押すと終了します")
print("\n【注意事項】")
print("- 深度画像はグレースケールまたはカラー画像どちらでも可")
print("- 出力画像は 'stereo_output_improved.jpg' として保存されます")
print("========================================\n")

print("深度表現方式:")
print("  A = 白が遠い/黒が近い")
print("  B = 黒が遠い/白が近い")
depth_mode = input("\n深度表現方式を選択 (A/B): ").upper()

# 画像ファイル選択
root = tk.Tk()
root.withdraw()

print("\nカラー画像を選択してください...")
color_path = filedialog.askopenfilename(title="カラー画像を選択")
if not color_path:
    print("カラー画像が選択されませんでした。終了します。")
    exit()

print("深度画像を選択してください...")
depth_path = filedialog.askopenfilename(title="深度画像を選択")
if not depth_path:
    print("深度画像が選択されませんでした。終了します。")
    exit()

root.destroy()

# 画像読み込み
print("\n画像を読み込み中...")
color_image = cv2.imread(color_path)
depth_image_raw = cv2.imread(depth_path)

if color_image is None:
    print("カラー画像の読み込みに失敗しました。")
    exit()
if depth_image_raw is None:
    print("深度画像の読み込みに失敗しました。")
    exit()

# 深度画像の明度抽出
if len(depth_image_raw.shape) == 3:
    depth_image = cv2.cvtColor(depth_image_raw, cv2.COLOR_BGR2GRAY)
else:
    depth_image = depth_image_raw

# メイン処理
target_height, target_width = color_image.shape[:2]
src_h, src_w = depth_image.shape[:2]

print(f"カラー画像サイズ: {target_width}x{target_height}")
print(f"深度画像サイズ: {src_w}x{src_h}")

# サイズ調整（切り抜き→パディング）
if src_h > target_height or src_w > target_width:
    start_y = max(0, (src_h - target_height) // 2)
    start_x = max(0, (src_w - target_width) // 2)
    end_y = min(src_h, start_y + target_height)
    end_x = min(src_w, start_x + target_width)
    depth_image = depth_image[start_y:end_y, start_x:end_x]

current_h, current_w = depth_image.shape[:2]
if current_h < target_height or current_w < target_width:
    pad_top = (target_height - current_h) // 2
    pad_bottom = target_height - current_h - pad_top
    pad_left = (target_width - current_w) // 2
    pad_right = target_width - current_w - pad_left

    depth_image = cv2.copyMakeBorder(
        depth_image, pad_top, pad_bottom, pad_left, pad_right,
        cv2.BORDER_REPLICATE
    )

# 深度値正規化と視差計算
print("\n視差を計算中...")
if depth_mode == 'A':
    # A: 白が遠い（255→DEPTH_MAX）、黒が近い（0→DEPTH_MIN）
    depth_values = DEPTH_MIN + (depth_image / 255.0) * (DEPTH_MAX - DEPTH_MIN)
else:
    # B: 黒が遠い（0→DEPTH_MAX）、白が近い（255→DEPTH_MIN）
    depth_values = DEPTH_MAX - (depth_image / 255.0) * (DEPTH_MAX - DEPTH_MIN)

depth_safe = np.clip(depth_values, DEPTH_MIN, DEPTH_MAX)
disparity = (FOCAL_LENGTH * BASELINE) / depth_safe
max_disparity = target_width * 0.15  # ステレオビジョンの一般的な範囲
disparity = np.clip(disparity, 0, max_disparity).astype(np.float32)

# 右画像生成
height, width = color_image.shape[:2]
right_image = np.zeros_like(color_image)
depth_buffer = np.full((height, width), np.inf)

print("右画像を生成中...")
for y in range(height):
    if y % 100 == 0:
        print(f"  進捗: {y}/{height} 行")
    for x in range(width):
        disp = disparity[y, x]
        new_x = int(x - disp)

        if 0 <= new_x < width:
            current_depth = depth_safe[y, x]

            if current_depth < depth_buffer[y, new_x]:
                right_image[y, new_x] = color_image[y, x]
                depth_buffer[y, new_x] = current_depth

print("\n穴埋め処理を実行中...")
right_image = improved_hole_filling(right_image, depth_image, disparity, max_disparity)

# 結果出力
stereo_pair = np.hstack([color_image, right_image])

output_filename = 'stereo_output_improved.jpg'
cv2.imwrite(output_filename, stereo_pair)
print(f"\n処理完了！")
print(f"ステレオペア画像を '{output_filename}' として保存しました")
print("左側：元画像、右側：右目用画像")
print("\n任意のキーを押すと終了します...")

cv2.imshow('ステレオペア画像', stereo_pair)
cv2.waitKey(0)
cv2.destroyAllWindows()