AnyDoorとは？

# PyTorch, Diffusers, Transformers, OpenCV, Scikit-imageなどの必須ライブラリをインストール
!pip install torch diffusers transformers opencv-python scikit-image

# 動画セグメンテーションの標準的なベンチマークであるDAVIS 2017データセットをダウンロード
!wget https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-trainval-480p.zip
# ダウンロードしたzipファイルを解凍し、DAVIS2017ディレクトリに展開
!unzip DAVIS-2017-trainval-480p.zip -d ./DAVIS2017

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Dinov2Model
from diffusers import UNet2DConditionModel, DDPMScheduler, AutoencoderKL
from torch.optim import AdamW
import cv2
import numpy as np
import random
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
import os
import glob
from tqdm import tqdm
from PIL import Image
import matplotlib.pyplot as plt
from skimage.metrics import peak_signal_noise_ratio as calculate_psnr
from skimage.metrics import structural_similarity as calculate_ssim


class MiniAnyDoor(nn.Module):
    def __init__(self, cross_attention_dim=768):
        super().__init__()
        # 各特徴抽出器を初期化
        self.hf_extractor = HFMapExtractor()
        self.id_extractor = IDExtractor(cross_attention_dim=cross_attention_dim)
        self.detail_extractor = DetailExtractor()

        # Stable Diffusion v1.5のU-Netをベースとしてロード
        self.unet = UNet2DConditionModel.from_pretrained(
            "runwayml/stable-diffusion-v1-5", subfolder="unet"
        )

        # U-Netのエンコーダ部分（down_blocks）の重みを凍結し、事前学習された知識を保持
        for name, param in self.unet.named_parameters():
            if "up_blocks" not in name:
                param.requires_grad = False

    def forward(self, target_img, target_mask, aligned_target, aligned_mask, scene_bg, shape_mask, noisy_latents, timesteps, drop_prob=0.1):
        # 詳細特徴の元となる高周波マップを生成
        hf_map = self.hf_extractor(aligned_target, aligned_mask)
        # シーン背景と高周波マップを合成してコラージュ画像を作成
        collage_img = scene_bg + hf_map

        # オブジェクトのID特徴と詳細特徴を抽出
        id_tokens = self.id_extractor(target_img * target_mask)
        detail_maps = self.detail_extractor(collage_img, shape_mask)

        # 学習時の正則化のため、一定確率で条件付けをドロップアウト
        if self.training and torch.rand(1).item() < drop_prob:
            id_tokens = torch.zeros_like(id_tokens) # IDトークンをゼロに
            detail_maps = [torch.zeros_like(d) for d in detail_maps] # 詳細マップをゼロに

        # U-Netにノイズのかかった潜在変数、タイムステップ、IDトークン、詳細マップを入力し、ノイズを予測
        noise_pred = self.unet(
            noisy_latents,
            timesteps,
            encoder_hidden_states=id_tokens,
            down_intrablock_additional_residuals=detail_maps, # 詳細マップをU-Netのデコーダに注入
        ).sample

        return noise_pred

class HFMapExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        # 水平方向と垂直方向のエッジを検出するためのSobelフィルタを定義
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32)
        # 畳み込み演算で使えるように、重みを(out_channels, in_channels, H, W)の形式に変形
        # register_bufferに登録することで、モデルのパラメータとしては扱われず、しかしstate_dictには保存される
        self.register_buffer('weight_x', sobel_x.view(1, 1, 3, 3))
        self.register_buffer('weight_y', sobel_y.view(1, 1, 3, 3))

    def forward(self, img_rgb, object_mask):
        # RGB画像をグレースケールに変換
        img_gray = 0.299 * img_rgb[:, 0:1, :, :] + 0.587 * img_rgb[:, 1:2, :, :] + 0.114 * img_rgb[:, 2:3, :, :]
        # Sobelフィルタを適用してx方向とy方向のエッジを計算
        edge_x = F.conv2d(img_gray, self.weight_x, padding=1)
        edge_y = F.conv2d(img_gray, self.weight_y, padding=1)
        # 両方向のエッジの強度を合成して、全体のエッジ（高周波成分）の強度を計算
        high_freq_intensity = torch.sqrt(edge_x**2 + edge_y**2)
        # 元のRGB画像にエッジ強度を乗算し、エッジ部分のみを強調した高周波マップを作成
        hf_rgb = img_rgb * high_freq_intensity
        # マスクをわずかに侵食させて、オブジェクトの輪郭付近の不要なアーティファクトを除去
        eroded_mask = 1.0 - F.max_pool2d(1.0 - object_mask, kernel_size=5, stride=1, padding=2)
        # 最終的な高周波マップにマスクを適用し、オブジェクト部分のみを抽出
        return hf_rgb * eroded_mask

class IDExtractor(nn.Module):
    def __init__(self, cross_attention_dim=768):
        super().__init__()
        # 事前学習済みの強力な自己教師ありモデルDINOv2をバックボーンとしてロード
        self.dino = Dinov2Model.from_pretrained("facebook/dinov2-base")

        # 特徴量をU-Netの埋め込み空間に射影するためのプロジェクターを定義
        hidden_size = self.dino.config.hidden_size
        # 論文に記載の通り、単純な線形層で特徴量を射影
        self.projector = nn.Linear(hidden_size, cross_attention_dim)

        # DINOv2のパラメータは学習中に更新しないように凍結
        for param in self.dino.parameters():
            param.requires_grad = False

    def forward(self, clean_target_img):
        # DINOv2は224x224の入力を想定しているため、リサイズ
        clean_target_img_224 = F.interpolate(clean_target_img, size=(224, 224), mode='bilinear', align_corners=False)
        # DINOv2で特徴を抽出
        outputs = self.dino(clean_target_img_224)
        # 最後の隠れ層の状態（パッチトークンとCLStoken）を取得
        tokens = outputs.last_hidden_state
        # プロジェクターを通して、U-Netのクロスアテンション層で扱える次元に変換
        id_tokens = self.projector(tokens)
        return id_tokens

class DetailExtractor(nn.Module):
    def __init__(self, in_channels=4):
        super().__init__()
        # 畳み込み層を使い、入力をU-Netの中間層と同じ解像度(64x64)までダウンサンプリング
        self.down_8x = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=3, stride=2, padding=1),  # 512 -> 256
            nn.SiLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),           # 256 -> 128
            nn.SiLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),          # 128 -> 64
            nn.SiLU()
        )

        # U-Netの各ブロックに対応する解像度の特徴マップを生成する層
        # 64x64解像度の特徴マップ (U-Netのup_blocks[0]に対応)
        self.feat1 = nn.Sequential(nn.Conv2d(128, 128, kernel_size=3, padding=1), nn.SiLU())
        # 32x32解像度の特徴マップ (U-Netのup_blocks[1]に対応)
        self.feat2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1), nn.SiLU())
        # 16x16解像度の特徴マップ (U-Netのup_blocks[2]に対応)
        self.feat3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1), nn.SiLU())

        # U-Netの各デコーダブロックの入力チャネル数(320, 640, 1280)に合わせるためのZero-Convolution層
        self.zero_convs = nn.ModuleList([
            nn.Conv2d(128, 320, kernel_size=1),
            nn.Conv2d(256, 640, kernel_size=1),
            nn.Conv2d(512, 1280, kernel_size=1)
        ])

        # Zero-Convolution層の重みとバイアスを0で初期化
        # これにより、学習初期では詳細情報がU-Netに影響を与えず、安定した学習が可能になる
        for m in self.zero_convs:
            nn.init.zeros_(m.weight)
            nn.init.zeros_(m.bias)

    def forward(self, collage_img, shape_mask):
        # 入力としてコラージュ画像と形状マスクをチャネル方向に連結
        x = torch.cat([collage_img, shape_mask], dim=1) # 入力サイズ: [B, 4, 512, 512]

        # 基準となる64x64の特徴マップを生成
        x_base = self.down_8x(x)

        # 階層的な特徴マップを抽出
        f1 = self.feat1(x_base)  # [B, 128, 64, 64]
        f2 = self.feat2(f1)      # [B, 256, 32, 32]
        f3 = self.feat3(f2)      # [B, 512, 16, 16]

        # Zero-Convolutionを適用して、U-Netデコーダの各ブロックのチャネル数に合わせる
        out1 = self.zero_convs[0](f1) # [B, 320, 64, 64]
        out2 = self.zero_convs[1](f2) # [B, 640, 32, 32]
        out3 = self.zero_convs[2](f3) # [B, 1280, 16, 16]

        # 3つの異なる解像度の特徴マップをリストとして返す
        return [out1, out2, out3]

class DavisAnyDoorDataset(Dataset):
    def __init__(self, davis_root, split="train", output_size=(512, 512), single_video=False):
        self.davis_root = davis_root
        self.output_size = output_size

        # 学習用(train)または検証用(val)の動画リストを読み込む
        imagesets_path = os.path.join(davis_root, "ImageSets", "2017", f"{split}.txt")
        with open(imagesets_path, 'r') as f:
            self.video_names = [line.strip() for line in f.readlines()]

        # デバッグや高速な検証のために、動画を1つに限定するオプション
        if single_video and len(self.video_names) > 0:
            self.video_names = [self.video_names[0]]

        # 画像とマスクが保存されているディレクトリのパスを設定
        self.img_dir = os.path.join(davis_root, "JPEGImages", "480p")
        self.mask_dir = os.path.join(davis_root, "Annotations", "480p")

        # 画像用の変換: Tensorへの変換と[-1, 1]への正規化
        self.transform = T.Compose([
            T.ToTensor(),
            T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])
        # マスク用の変換: Tensorへの変換のみ
        self.mask_transform = T.ToTensor()

    def __len__(self):
        # データセット内の動画の総数を返す
        return len(self.video_names)

    def _perturb_mask(self, mask):
        # データ拡張: マスクにランダムな変形（膨張・侵食など）を加えて、形状の多様性を学習させる
        if random.random() < 0.3:
            y, x = np.where(mask > 0)
            if len(y) > 0 and len(x) > 0:
                y1, y2, x1, x2 = y.min(), y.max(), x.min(), x.max()
                box_mask = np.zeros_like(mask)
                box_mask[y1:y2, x1:x2] = 1.0
                return box_mask

        kernel_size = random.randint(3, 15)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)

        if random.random() < 0.5:
            perturbed = cv2.dilate(mask, kernel, iterations=1)
        else:
            perturbed = cv2.erode(mask, kernel, iterations=1)

        return perturbed

    def _get_bbox(self, mask):
        # マスクからバウンディングボックス（オブジェクトの矩形領域）の座標を取得
        y, x = np.where(mask > 0)
        if len(y) == 0 or len(x) == 0:
            return 0, 0, mask.shape[1], mask.shape[0]
        return x.min(), y.min(), x.max(), y.max()

    def _align_image(self, src_img, src_mask, dst_mask):
        # ソースオブジェクトを、宛先マスクのバウンディングボックスに合わせてリサイズ＆ペーストする
        x1_s, y1_s, x2_s, y2_s = self._get_bbox(src_mask)
        x1_d, y1_d, x2_d, y2_d = self._get_bbox(dst_mask)

        aligned_img = np.zeros_like(src_img)
        aligned_mask = np.zeros_like(src_mask)

        w_d, h_d = x2_d - x1_d, y2_d - y1_d

        if w_d > 0 and h_d > 0 and (x2_s - x1_s) > 0 and (y2_s - y1_s) > 0:
            crop_img = src_img[y1_s:y2_s, x1_s:x2_s]
            crop_mask = src_mask[y1_s:y2_s, x1_s:x2_s]

            aligned_img[y1_d:y2_d, x1_d:x2_d] = cv2.resize(crop_img, (w_d, h_d))
            aligned_mask[y1_d:y2_d, x1_d:x2_d] = cv2.resize(crop_mask, (w_d, h_d), interpolation=cv2.INTER_NEAREST)

        return aligned_img, aligned_mask

    def __getitem__(self, idx):
        # 1つのデータサンプルを生成するメインの関数
        video_name = self.video_names[idx]
        img_paths = sorted(glob.glob(os.path.join(self.img_dir, video_name, "*.jpg")))
        mask_paths = sorted(glob.glob(os.path.join(self.mask_dir, video_name, "*.png")))

        num_frames = len(img_paths)
        # 1つの動画からランダムに2つのフレーム（i, j）を選択
        frame_i_idx, frame_j_idx = random.sample(range(num_frames), 2)

        # 画像とマスクを読み込む
        img_i = cv2.imread(img_paths[frame_i_idx])[..., ::-1] # OpenCV(BGR)からRGBに変換
        img_j = cv2.imread(img_paths[frame_j_idx])[..., ::-1]

        mask_i_raw = cv2.imread(mask_paths[frame_i_idx], cv2.IMREAD_GRAYSCALE)
        mask_j_raw = cv2.imread(mask_paths[frame_j_idx], cv2.IMREAD_GRAYSCALE)
        # マスクを0-1のfloat値に変換
        mask_i = (mask_i_raw > 0).astype(np.float32)
        mask_j = (mask_j_raw > 0).astype(np.float32)

        # 画像とマスクを指定された出力サイズにリサイズ
        img_i = cv2.resize(img_i, self.output_size, interpolation=cv2.INTER_LINEAR)
        img_j = cv2.resize(img_j, self.output_size, interpolation=cv2.INTER_LINEAR)
        mask_i = cv2.resize(mask_i, self.output_size, interpolation=cv2.INTER_NEAREST)
        mask_j = cv2.resize(mask_j, self.output_size, interpolation=cv2.INTER_NEAREST)

        # ターゲットオブジェクト画像（フレームiのオブジェクト部分）を作成
        target_img = (img_i * mask_i[:, :, None]).astype(np.uint8)

        # 形状マスク（フレームjのマスクに変形を加えたもの）を生成
        shape_mask = self._perturb_mask(mask_j)

        # 正解画像（フレームj）と、シーン背景（フレームjの背景部分）を作成
        gt_image = img_j
        scene_bg = (img_j * (1.0 - mask_j[:, :, None])).astype(np.uint8)

        # モデルの入力となる各種データを辞書形式で返す
        return {
            "gt_image": self.transform(gt_image),
            "target_img": self.transform(target_img),
            "target_mask": self.mask_transform(mask_i),
            "aligned_target": self.transform(target_img), # 本来は位置合わせした画像
            "aligned_mask": self.mask_transform(mask_i),   # 本来は位置合わせしたマスク
            "scene_bg": self.transform(scene_bg),
            "shape_mask": self.mask_transform(shape_mask),
            "modality": "video" # データの種類を示す
        }

def sample_timesteps(batch_size, modality, num_train_timesteps=1000, device="cpu"):
    # modalityが'video'の場合
    if modality == "video":
        # 50%の確率で、ノイズ除去プロセスの前半（500-1000）からタイムステップをサンプリング
        # 動画データは構造やポーズの学習に適しているため、初期ステップを重視
        bias_probs = torch.rand(batch_size, device=device)
        timesteps = torch.where(
            bias_probs < 0.5,
            torch.randint(num_train_timesteps // 2, num_train_timesteps, (batch_size,), device=device),
            torch.randint(0, num_train_timesteps, (batch_size,), device=device)
        )
    # modalityが'image'の場合（今回は使用しないが、参考として記載）
    elif modality == "image":
        # 50%の確率で、ノイズ除去プロセスの後半（0-500）からタイムステップをサンプリング
        # 静止画データはディテールの学習に適しているため、後期ステップを重視
        bias_probs = torch.rand(batch_size, device=device)
        timesteps = torch.where(
            bias_probs < 0.5,
            torch.randint(0, num_train_timesteps // 2, (batch_size,), device=device),
            torch.randint(0, num_train_timesteps, (batch_size,), device=device)
        )
    # 上記以外の場合は、全範囲から均等にサンプリング
    else:
        timesteps = torch.randint(0, num_train_timesteps, (batch_size,), device=device)
    return timesteps

def generate_image(model, vae, scheduler, batch, device, num_inference_steps=50, guidance_scale=1.0):
    model.eval() # モデルを評価モードに設定
    with torch.no_grad(): # 勾配計算を無効化
        # 推論に必要なデータをデバイスに転送
        target_img = batch["target_img"].to(device)
        target_mask = batch["target_mask"].to(device)
        aligned_target = batch["aligned_target"].to(device)
        aligned_mask = batch["aligned_mask"].to(device)
        scene_bg = batch["scene_bg"].to(device)
        shape_mask = batch["shape_mask"].to(device)

        B = target_img.shape[0] # バッチサイズ

        # --- 1. 条件付き (Conditional) 特徴量の抽出 ---
        # IDトークンと詳細マップをモデルから抽出
        hf_map = model.hf_extractor(aligned_target, aligned_mask)
        collage_img = scene_bg + hf_map
        id_tokens_cond = model.id_extractor(target_img * target_mask)
        detail_maps_cond = model.detail_extractor(collage_img, shape_mask)

        # --- 2. 無条件 (Unconditional) 特徴量の生成 ---
        # CFGのため、条件付けを行わない場合のダミー特徴量（ゼロベクトル）を用意
        id_tokens_uncond = torch.zeros_like(id_tokens_cond)
        detail_maps_uncond = [torch.zeros_like(d) for d in detail_maps_cond]

        # --- 3. 特徴量の結合 ---
        # バッチ方向に無条件と条件付きの特徴量を結合
        # [uncond, cond] の順で並べる
        id_tokens_both = torch.cat([id_tokens_uncond, id_tokens_cond])
        detail_maps_both = [torch.cat([u, c]) for u, c in zip(detail_maps_uncond, detail_maps_cond)]

        # --- 4. ノイズ除去ループ ---
        # 純粋なノイズから潜在変数を初期化
        latents = torch.randn((B, 4, 64, 64), device=device)
        # スケジューラに推論ステップ数を設定
        scheduler.set_timesteps(num_inference_steps)

        for t in tqdm(scheduler.timesteps, desc="Denoising Steps"):
            # U-Netへの入力も無条件・条件付きのペアにするため、潜在変数を2倍に複製
            latent_model_input = torch.cat([latents, latents])

            # スケジューラによっては入力のスケーリングが必要
            if hasattr(scheduler, "scale_model_input"):
                latent_model_input = scheduler.scale_model_input(latent_model_input, t)

            # 1回のU-Netフォワードパスで、無条件と条件付きの両方のノイズ予測を同時に計算
            noise_pred = model.unet(
                latent_model_input,
                t,
                encoder_hidden_states=id_tokens_both,
                down_intrablock_additional_residuals=detail_maps_both,
            ).sample

            # --- 5. CFGの適用 ---
            # 予測結果を無条件(uncond)と条件付き(cond)に分割
            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)

            # CFGの計算式: noise = uncond_noise + guidance_scale * (cond_noise - uncond_noise)
            # 条件の方向へ、guidance_scale分だけ強く予測を誘導する
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)

            # スケジューラを使って1ステップ前の潜在変数を計算
            latents = scheduler.step(noise_pred, t, latents).prev_sample

        # --- 6. VAEデコーダによる画像化 ---
        # 潜在変数をスケーリングし、VAEデコーダでピクセル空間の画像に変換
        latents = (1 / vae.config.scaling_factor) * latents
        generated_image = vae.decode(latents).sample

        # 画像を[0, 1]の範囲にクリップし、Numpy配列(0-255のuint8)に変換
        generated_image = (generated_image / 2 + 0.5).clamp(0, 1)
        generated_image = generated_image.cpu().permute(0, 2, 3, 1).numpy()
        generated_image = (generated_image * 255).astype(np.uint8)

    return generated_image[0]

if __name__ == "__main__":
    # GPUが利用可能ならGPUを、そうでなければCPUを使用
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # モデル、VAE、スケジューラを初期化し、指定したデバイスに配置
    model = MiniAnyDoor().to(device)
    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae").to(device)
    noise_scheduler = DDPMScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
    # VAEは特徴抽出にのみ使い、学習は行わないため勾配計算を無効化
    vae.requires_grad_(False)

    # オプティマイザを定義。学習対象のパラメータ（requires_grad=Trueのもの）のみを渡す
    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)

    def train_step(batch):
        model.train() # モデルを学習モードに設定
        optimizer.zero_grad() # 勾配をリセット

        # バッチから各種データを取得し、デバイスに転送
        gt_image = batch["gt_image"].to(device)
        target_img = batch["target_img"].to(device)
        target_mask = batch["target_mask"].to(device)
        scene_bg = batch["scene_bg"].to(device)
        shape_mask = batch["shape_mask"].to(device)
        modality = batch["modality"][0]

        B = gt_image.shape[0] # バッチサイズ

        with torch.no_grad(): # このブロック内では勾配を計算しない
            # VAEエンコーダで正解画像を潜在空間に変換
            latents = vae.encode(gt_image).latent_dist.sample()
            # VAEのスケーリングファクターを適用
            latents = latents * vae.config.scaling_factor

        # 純粋なノイズを生成
        noise = torch.randn_like(latents)
        # 適応的サンプリングでタイムステップを決定
        timesteps = sample_timesteps(B, modality, num_train_timesteps=1000, device=device)
        # 元の潜在変数にノイズを付加
        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

        # モデルに入力し、ノイズを予測
        noise_pred = model(
            target_img=target_img,
            target_mask=target_mask,
            aligned_target=batch["aligned_target"].to(device),
            aligned_mask=batch["aligned_mask"].to(device),
            scene_bg=scene_bg,
            shape_mask=shape_mask,
            noisy_latents=noisy_latents,
            timesteps=timesteps
        )

        # 予測したノイズ(noise_pred)と元のノイズ(noise)の平均二乗誤差(MSE)を損失として計算
        loss = F.mse_loss(noise_pred, noise)
        # 損失に基づいて勾配を計算
        loss.backward()
        # オプティマイザでモデルのパラメータを更新
        optimizer.step()
        return loss.item()

    # DAVISデータセットのルートディレクトリを指定
    davis_root = "./DAVIS2017/DAVIS"

    # データセットが存在するか確認
    if os.path.exists(davis_root):
        print("DAVIS2017 dataset found. Loading SINGLE video for quick verification...")
        # 動作検証のため、1つの動画のみを含むデータセットをロード
        dataset = DavisAnyDoorDataset(davis_root, split="train", single_video=True)
        # データローダーを作成。バッチサイズ5でデータを供給
        dataloader = DataLoader(dataset, batch_size=5, shuffle=True)

        # 学習のエポック数を設定
        num_epochs = 1500

        print("\n--- 学習フェーズ開始 (1つの動画への過学習) ---")
        # 指定されたエポック数だけ学習ループを実行
        for epoch in range(num_epochs):
            epoch_loss = 0 # エポックごとの損失を記録する変数
            # データローダーからバッチ単位でデータを取得
            for step, batch in enumerate(dataloader):
                # 1ステップ分の学習を実行し、損失を取得
                loss = train_step(batch)
                epoch_loss += loss
            # 10エポックごとに進捗（損失）を出力
            if (epoch+1) % 10 == 0:
                print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}")
        print("--- 学習フェーズ終了 ---\n")

        print("--- 推論フェーズ開始 ---")

        # 学習に使った動画の名前を取得
        video_name = dataset.video_names[0]
        img_paths = sorted(glob.glob(os.path.join(dataset.img_dir, video_name, "*.jpg")))
        mask_paths = sorted(glob.glob(os.path.join(dataset.mask_dir, video_name, "*.png")))

        # ソースオブジェクト(A)、ターゲットシーン(B)、比較用(C)のフレームインデックスを決定
        idx_A = 0
        idx_C = len(img_paths) - 1
        idx_B = len(img_paths) // 2

        # 指定したインデックスの画像とマスクをロードするヘルパー関数
        def load_img_mask(idx):
            img = cv2.imread(img_paths[idx])[..., ::-1]
            mask = cv2.imread(mask_paths[idx], cv2.IMREAD_GRAYSCALE)
            mask = (mask > 0).astype(np.float32)
            img = cv2.resize(img, dataset.output_size, interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, dataset.output_size, interpolation=cv2.INTER_NEAREST)
            return img, mask

        # 各フレームの画像とマスクをロード
        img_A, mask_A = load_img_mask(idx_A) # ソースオブジェクト
        img_C, mask_C = load_img_mask(idx_C) # 比較用
        img_B, mask_B = load_img_mask(idx_B) # ターゲットシーン / 正解画像

        # 推論用の入力データを作成
        aligned_target_np, aligned_mask_np = dataset._align_image(img_A, mask_A, mask_B)
        target_img_np = (img_A * mask_A[:, :, None]).astype(np.uint8)
        scene_bg_np = (img_B * (1.0 - mask_B[:, :, None])).astype(np.uint8)

        # PyTorchのTensorに変換し、バッチ次元を追加
        inference_batch = {
            "target_img": dataset.transform(target_img_np).unsqueeze(0),
            "target_mask": dataset.mask_transform(mask_A).unsqueeze(0),
            "aligned_target": dataset.transform(target_img_np).unsqueeze(0),
            "aligned_mask": dataset.mask_transform(mask_A).unsqueeze(0),
            "scene_bg": dataset.transform(scene_bg_np).unsqueeze(0),
            "shape_mask": dataset.mask_transform(mask_B).unsqueeze(0)
        }

        # 画像生成関数を呼び出して推論を実行
        result_img_array = generate_image(
            model=model,
            vae=vae,
            scheduler=noise_scheduler,
            batch=inference_batch,
            device=device,
            num_inference_steps=50
        )

        # --- 精度指標 (PSNR / SSIM) の計算 ---
        # 正解画像 (img_B) と生成画像 (result_img_array) を比較
        # PSNR (Peak Signal-to-Noise Ratio): 画質の劣化を示す指標。高いほど良い。
        psnr_value = calculate_psnr(img_B, result_img_array, data_range=255)
        # SSIM (Structural Similarity Index): 構造的な類似度を示す指標。1.0に近いほど良い。
        ssim_value = calculate_ssim(img_B, result_img_array, data_range=255, channel_axis=-1)

        print("\n--- 評価結果 ---")
        print(f"PSNR : {psnr_value:.2f} dB (高いほど良い)")
        print(f"SSIM : {ssim_value:.4f} (1.0に近いほど良い)")
        print("----------------\n")

        # --- Matplotlibによるプレビュー表示 ---
        fig, axes = plt.subplots(1, 4, figsize=(20, 5))

        axes[0].imshow(img_A)
        axes[0].set_title(f"Frame A (Source Object)\nIdx: {idx_A}")

        axes[1].imshow(img_C)
        axes[1].set_title(f"Frame C (Distant)\nIdx: {idx_C}")

        axes[2].imshow(img_B)
        axes[2].set_title(f"Frame B (Ground Truth)\nIdx: {idx_B}")

        axes[3].imshow(result_img_array)
        axes[3].set_title(f"Generated Frame B\nPSNR: {psnr_value:.2f} | SSIM: {ssim_value:.4f}")

        for ax in axes:
            ax.axis('off')

        plt.tight_layout()
        plt.show()

    else:
        print(f"Dataset directory '{davis_root}' not found. Please check the path.")

DAVIS2017 dataset found. Loading SINGLE video for quick verification...

--- 学習フェーズ開始 (1つの動画への過学習) ---
Epoch [10/1500], Loss: 0.0735
Epoch [20/1500], Loss: 0.0661
Epoch [30/1500], Loss: 0.0401
...
Epoch [1480/1500], Loss: 0.0593
Epoch [1490/1500], Loss: 0.0402
Epoch [1500/1500], Loss: 0.0085
--- 学習フェーズ終了 ---

--- 推論フェーズ開始 ---
--- 評価結果 ---
PSNR : 15.95 dB (高いほど良い)
SSIM : 0.1999 (1.0に近いほど良い)
----------------

import os

os.chdir("/content/AnyDoor")

pth_path = "pretrained_models/dinov2_vitg14_pretrain.pth"

yaml_path = "configs/anydoor.yaml"
with open(yaml_path, "r", encoding="utf-8") as f:
    content = f.read()

content = content.replace("path/dinov2_vitg14_pretrain.pth", pth_path)

with open(yaml_path, "w", encoding="utf-8") as f:
    f.write(content)

import cv2
import torch
import numpy as np
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection, SamModel, SamProcessor
import os
import sys

# Google Colabの環境依存エラー（einops関連）を回避するためのモジュールリロード処理
for module_name in list(sys.modules.keys()):
    if "einops" in module_name:
        del sys.modules[module_name]

# Python 3.12以降で廃止されたImpImporterをエミュレートするパッチ（古いライブラリとの互換性維持）
import pkgutil
if not hasattr(pkgutil, 'ImpImporter'):
    class ImpImporter:
        pass
    pkgutil.ImpImporter = ImpImporter

# AnyDoor公式のヘルパー関数やモデル定義をインポート
from predict import process_pairs
from cldm.model import create_model, load_state_dict
from pytorch_lightning import seed_everything

# バックエンドにPyTorchを指定
os.environ["EINOPS_BACKEND"] = "torch"
from cldm.ddim_hacked import DDIMSampler
import matplotlib.pyplot as plt

# 実行デバイス（GPUが利用可能ならCUDA、なければCPU）を決定
device = "cuda" if torch.cuda.is_available() else "cpu"

# ==============================================================================
# --- 0. GroundingDINO & SAM（セグメンテーションモデル）のセットアップ ---
# ==============================================================================
print("GroundingDINOとSAMのモデルをロード中...")
# 物体検出モデル GroundingDINO のプロセッサとモデルをロード
gd_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
gd_model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base").to(device)

# 高精度セグメンテーションモデル SAM のプロセッサとモデルをロード
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)

# テキスト指示から対象のマスク画像を自動生成し、検出枠を視覚化して保存する関数
def generate_mask_and_visualize(image_path, text_prompt, output_mask_path, output_vis_path):
    # 画像をRGB形式で読み込む
    image = Image.open(image_path).convert("RGB")
    
    # GroundingDINOによる物体検出を実行
    inputs = gd_processor(images=image, text=text_prompt, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = gd_model(**inputs)
    
    # 検出結果のポストプロセス（境界ボックスの抽出）
    results = gd_processor.post_process_grounded_object_detection(
        outputs, inputs.input_ids, threshold=0.15, text_threshold=0.25, target_sizes=[image.size[::-1]]
    )[0]
    
    # オブジェクトが検出されなかった場合はエラー
    if len(results["boxes"]) == 0:
        raise ValueError(f"エラー: '{text_prompt}' が検出できませんでした。")
        
    # 最もスコアの高い境界ボックスを取得
    best_idx = torch.argmax(results["scores"])
    bbox = results["boxes"][best_idx].tolist()

    # 境界ボックス（物体検出枠）を画像に描画してデバッグ用に保存
    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    x_min, y_min, x_max, y_max = map(int, bbox)
    cv2.rectangle(img_cv, (x_min, y_min), (x_max, y_max), (0, 0, 255), 3)
    cv2.putText(img_cv, text_prompt, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
    cv2.imwrite(output_vis_path, img_cv)

    # 検出された境界ボックスをプロンプトとして、SAMにより高精度なオブジェクトマスクを生成
    inputs_sam = sam_processor(image, input_boxes=[[bbox]], return_tensors="pt").to(device)
    with torch.no_grad():
        outputs_sam = sam_model(**inputs_sam)
    
    # マスク画像のポストプロセス
    mask = sam_processor.image_processor.post_process_masks(
        outputs_sam.pred_masks.cpu(), inputs_sam["original_sizes"].cpu(), inputs_sam["reshaped_input_sizes"].cpu()
    )[0][0][0].numpy() 
    
    # マスク画像を0-255の白黒画像（2値マスク）として保存
    mask_img = (mask * 255).astype(np.uint8)
    cv2.imwrite(output_mask_path, mask_img)
    return output_mask_path

# ==============================================================================
# --- 1. マスク生成の実行 ---
# ==============================================================================
ref_image_path = "orange.jpg"      # ソース画像（移植したいオブジェクトを含む画像）
tar_image_path = "apple.jpg"       # ターゲット画像（移植先の背景となる画像）

ref_prompt = "orange."             # 検出したいソースオブジェクトのプロンプト
tar_prompt = "apple"               # 検出したいターゲット位置（入れ替え対象）のプロンプト

print("\n--- ソースオブジェクトの処理 ---")
ref_mask_path = "mask.png"
# ソース画像のオブジェクトを検出してマスクを作成
generate_mask_and_visualize(ref_image_path, ref_prompt, ref_mask_path, "ref_bbox_vis.jpg")

print("\n--- ターゲットシーンの処理 ---")
tar_mask_path = "taget_mask.png"
# ターゲット画像内のオブジェクトを検出してマスクを作成
generate_mask_and_visualize(tar_image_path, tar_prompt, tar_mask_path, "tar_bbox_vis.jpg")

# ==============================================================================
# --- 2. AnyDoorモデルによる前処理と画像合成（テレポート推論） ---
# ==============================================================================
model_ckpt = "pretrained_models/anydoor.ckpt" # ダウンロードしたAnyDoorの学習済み重み
config_file = "configs/anydoor.yaml"           # 設定ファイルのパス

print("\nAnyDoorのモデルをロード中...")
# 設定ファイルに基づきAnyDoorモデルを定義
model = create_model(config_file).cpu()
# 学習済みの重みをGPU（CUDA）にロード
model.load_state_dict(load_state_dict(model_ckpt, location='cuda'))
model = model.cuda()
model.eval()
# 再現性確保のため乱数シードを固定
seed_everything(42)

# OpenCVを使って画像とセグメンテーションマスクを読み込む
ref_image = cv2.imread(ref_image_path)
tar_image = cv2.imread(tar_image_path)
ref_mask = cv2.imread(ref_mask_path, cv2.IMREAD_GRAYSCALE)
tar_mask = cv2.imread(tar_mask_path, cv2.IMREAD_GRAYSCALE)

# ノイズやブレを防止するため、マスク画像を確実に2値化（0または255）
_, ref_mask = cv2.threshold(ref_mask, 127, 255, cv2.THRESH_BINARY)
_, tar_mask = cv2.threshold(tar_mask, 127, 255, cv2.THRESH_BINARY)

# AnyDoor公式の前処理関数（process_pairs）を実行してモデルの入力データを生成
# 形状コントロールを無効化（enable_shape_control=False）することで、オブジェクトの元の形状を維持して合成
item = process_pairs(ref_image, ref_mask, tar_image, tar_mask, enable_shape_control=False)

print("AnyDoorのサンプラーを準備中...")
# DDIMサンプラーを初期化
sampler = DDIMSampler(model)

# モデルに入力するためのバッチデータを準備（Numpy配列からPyTorchのTensorに変換してGPUへ転送）
batch = {}
for k, v in item.items():
    if isinstance(v, np.ndarray):
        tensor_v = torch.from_numpy(v).float()
        if tensor_v.ndim == 3:
            # 軸を (H, W, C) から PyTorch標準の (C, H, W) に並べ替え
            tensor_v = tensor_v.permute(2, 0, 1)
        elif tensor_v.ndim == 2:
            tensor_v = tensor_v.unsqueeze(0)
        # バッチ次元（B=1）を追加してデバイスに配置
        batch[k] = tensor_v.unsqueeze(0).to(device)

print("AnyDoorの推論を実行中...")
with torch.no_grad():
    ref_img, hint_tensor = batch['ref'], batch['hint']
    
    # ソース画像からオブジェクトのID特徴を抽出
    c_crossattn = model.get_learned_conditioning(ref_img)
    # 抽出したID特徴と詳細マップ（hint_tensor）を条件（cond）として指定
    cond = {"c_crossattn": [c_crossattn], "c_concat": [hint_tensor]}
    
    # Classifier-Free Guidance (CFG) のための無条件（Unconditional）のダミー条件を準備
    uc_cross = model.get_unconditional_conditioning(1)
    uc = {"c_crossattn": [uc_cross], "c_concat": [hint_tensor]}

    # DDIMサンプラーを用いて50ステップのサンプリング（ノイズ除去ループ）を実行
    # guidance_scale（強さ）には、公式推奨の 4.5 を指定
    samples, _ = sampler.sample(
        S=50, batch_size=1, shape=(4, 64, 64),
        conditioning=cond, verbose=False,
        unconditional_guidance_scale=4.5, unconditional_conditioning=uc
    )

    # 潜在表現（サンプル）をファーストステージのデコーダを通してピクセル空間の画像に復元
    x_samples = model.decode_first_stage(samples)
    # 画像のテンソル値を [-1, 1] から [0, 255] の範囲に正規化・スケーリング
    x_samples = (torch.clamp(x_samples, -1.0, 1.0) + 1.0) / 2.0 * 255.0
    # Numpy配列に変換し、次元を (C, H, W) から (H, W, C) に戻す
    out_image = x_samples.cpu().numpy()[0].transpose(1, 2, 0).astype(np.uint8)

# 出力画像（RGB形式）を、OpenCVでの保存用にBGR形式に変換してからファイルに保存
cv2.imwrite("output_anydoor.jpg", cv2.cvtColor(out_image, cv2.COLOR_RGB2BGR))
print("推論が完了しました！ 'output_anydoor.jpg' を確認してください！")

# 合成画像のプレビューを表示
plt.figure(figsize=(8, 8))
plt.imshow(out_image)
plt.axis("off")  # 軸の目盛りを非表示にする
plt.title("AnyDoor Output")
plt.show()