{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# depth2개비교.ipynb (또는 .py 스크립트 형태로 사용 가능)\n",
    "# - Distill-Any-Depth, Depth-Anything-V2 비교\n",
    "# - ModuleNotFoundError 해결: sys.path.append(...) 추가\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import shutil\n",
    "import subprocess\n",
    "import numpy as np\n",
    "import cv2\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "# -------------------------------------------------------------\n",
    "# 0) 경로 설정\n",
    "# -------------------------------------------------------------\n",
    "# Distill-Any-Depth 관련 경로\n",
    "DISTILL_BASE_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/모델들/Distill-Any-Depth\"\n",
    "DISTILL_INFER_SCRIPT = os.path.join(DISTILL_BASE_DIR, \"scripts/00_infer.sh\")\n",
    "DISTILL_INPUT_DIR = os.path.join(DISTILL_BASE_DIR, \"data\", \"input\")\n",
    "DISTILL_OUTPUT_DIR = os.path.join(DISTILL_BASE_DIR, \"output\", \"large\")\n",
    "\n",
    "# Depth-Anything-V2 관련 경로\n",
    "DEPTH_ANYTHING_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/회사용/4090_옮길거/workspace/Depth-Anything-V2\"\n",
    "DEPTH_ANYTHING_CHECKPOINT_VITL = os.path.join(DEPTH_ANYTHING_DIR, \"checkpoints/depth_anything_v2_vitl.pth\")\n",
    "DEPTH_ENCODER_TYPE = \"vitl\"\n",
    "\n",
    "# 입력 이미지 폴더 (Distill-Any-Depth와 동일 폴더로 가정)\n",
    "INPUT_FOLDER = DISTILL_INPUT_DIR  # \"/.../Distill-Any-Depth/data/input\"\n",
    "# 최종 비교 결과 저장 폴더\n",
    "COMPARE_DEPTH_SAVE_DIR = \"/ltb/media/ltb/90887173887158A4/Users/ltb/compare_depth\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [1] run_distill_any_depth (수정됨)\n",
    "def run_distill_any_depth():\n",
    "    \"\"\"\n",
    "    Distill-Any-Depth로 data/input 폴더 전체 이미지를 추론 -> output/large/ 아래에 결과 생성\n",
    "    이후, image_logs 폴더에 생성된 이미지를 compare_depth 폴더로 복사하여\n",
    "    distill_ 접두어를 붙여 저장.\n",
    "    \"\"\"\n",
    "    import shutil\n",
    "    import os\n",
    "    import subprocess\n",
    "    \n",
    "    print(\"[Distill-Any-Depth] 전체 이미지 추론을 시작합니다.\")\n",
    "    \n",
    "    # Distill output 폴더 정리\n",
    "    if os.path.exists(DISTILL_OUTPUT_DIR):\n",
    "        for old_file in os.listdir(DISTILL_OUTPUT_DIR):\n",
    "            path_to_remove = os.path.join(DISTILL_OUTPUT_DIR, old_file)\n",
    "            if os.path.isfile(path_to_remove):\n",
    "                os.remove(path_to_remove)\n",
    "            else:\n",
    "                shutil.rmtree(path_to_remove, ignore_errors=True)\n",
    "    else:\n",
    "        os.makedirs(DISTILL_OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "    # subprocess로 00_infer.sh 실행\n",
    "    prev_cwd = os.getcwd()\n",
    "    os.chdir(DISTILL_BASE_DIR)\n",
    "    try:\n",
    "        cmd = f\"bash {DISTILL_INFER_SCRIPT}\"\n",
    "        subprocess.run(cmd, shell=True, check=True)\n",
    "    except subprocess.CalledProcessError as e:\n",
    "        print(\"[Distill-Any-Depth] 추론 에러 발생:\", e)\n",
    "    finally:\n",
    "        os.chdir(prev_cwd)\n",
    "\n",
    "    print(\"[Distill-Any-Depth] 추론 완료!\")\n",
    "\n",
    "    # Distill 결과가 /output/large/image_logs 폴더에 생성된다고 가정\n",
    "    distill_image_logs = os.path.join(DISTILL_OUTPUT_DIR, \"image_logs\")\n",
    "    if not os.path.exists(distill_image_logs):\n",
    "        print(f\"[Distill-Any-Depth] image_logs 폴더가 없습니다: {distill_image_logs}\")\n",
    "    else:\n",
    "        # compare_depth 폴더에 복사\n",
    "        os.makedirs(COMPARE_DEPTH_SAVE_DIR, exist_ok=True)\n",
    "        distill_logs = os.listdir(distill_image_logs)\n",
    "        for f in distill_logs:\n",
    "            src_path = os.path.join(distill_image_logs, f)\n",
    "            if os.path.isfile(src_path):\n",
    "                # distill_ 접두어를 붙인 새 이름\n",
    "                dst_name = f\"distill_{f}\"\n",
    "                dst_path = os.path.join(COMPARE_DEPTH_SAVE_DIR, dst_name)\n",
    "                shutil.copy2(src_path, dst_path)\n",
    "                print(f\"[Distill-Any-Depth] Copied => {os.path.abspath(dst_path)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 2) Depth-Anything-V2 모델 로드\n",
    "# -------------------------------------------------------------\n",
    "def load_depth_anything_v2(checkpoint_path, device=\"cuda\"):\n",
    "    \"\"\"\n",
    "    Depth-Anything-V2 모델(vitl) 불러오기\n",
    "    - sys.path.append(DEPTH_ANYTHING_DIR)를 통해 depth_anything_v2.dpt 모듈이 인식되도록 함.\n",
    "    \"\"\"\n",
    "    # 핵심: sys.path.append(...) 로 검색 경로에 추가\n",
    "    if DEPTH_ANYTHING_DIR not in sys.path:\n",
    "        sys.path.append(DEPTH_ANYTHING_DIR)\n",
    "\n",
    "    # 이제 import 가능\n",
    "    from depth_anything_v2.dpt import DepthAnythingV2\n",
    "\n",
    "    model = DepthAnythingV2(\n",
    "        encoder=DEPTH_ENCODER_TYPE,\n",
    "        features=256,\n",
    "        out_channels=[256, 512, 1024, 1024]\n",
    "    )\n",
    "    model.load_state_dict(torch.load(checkpoint_path, map_location=device))\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 3) Depth-Anything-V2로 이미지들 처리\n",
    "# -------------------------------------------------------------\n",
    "def run_depth_anything_on_images(model, input_images, device=\"cuda\", resize_size=700):\n",
    "    \"\"\"\n",
    "    input_images 리스트에 있는 모든 이미지를 불러와\n",
    "    Depth-Anything-V2로 inference 후 depth map(배열)을 리턴\n",
    "    return: { \"image_filename\": depth_map(np array), ... }\n",
    "    \"\"\"\n",
    "    depth_results = {}\n",
    "    with torch.no_grad():\n",
    "        for img_path in input_images:\n",
    "            img_name = os.path.basename(img_path)\n",
    "            pil_img = Image.open(img_path).convert('RGB')\n",
    "            img_np = np.array(pil_img)\n",
    "            depth_map = model.infer_image(img_np, resize_size)\n",
    "            depth_map = depth_map.squeeze()  # (H, W)\n",
    "            depth_results[img_name] = depth_map\n",
    "    return depth_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 4) 유틸리티 함수: 정규화 및 컬러맵 변환\n",
    "# -------------------------------------------------------------\n",
    "def normalize_minmax(depth_map):\n",
    "    d_min, d_max = depth_map.min(), depth_map.max()\n",
    "    if d_max - d_min < 1e-6:\n",
    "        return np.zeros_like(depth_map, dtype=np.float32)\n",
    "    return (depth_map - d_min) / (d_max - d_min)\n",
    "\n",
    "def depth_to_colormap(depth_map, cmap_name=\"plasma\"):\n",
    "    depth_map_norm = normalize_minmax(depth_map)\n",
    "    cmap = plt.get_cmap(cmap_name)\n",
    "    depth_color = cmap(depth_map_norm)[:, :, :3]  # RGBA 중 A는 제외\n",
    "    depth_color = (depth_color * 255).astype(np.uint8)\n",
    "    return depth_color\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 5) 비교 이미지 만들기\n",
    "# -------------------------------------------------------------\n",
    "def create_comparison_image(original_img, depth_img_1, depth_img_2, label1=\"ModelA\", label2=\"ModelB\"):\n",
    "    \"\"\"\n",
    "    original_img: (H, W, 3)\n",
    "    depth_img_1, depth_img_2: (H, W, 3)\n",
    "    \"\"\"\n",
    "    h, w, _ = original_img.shape\n",
    "    concat_img = np.concatenate([original_img, depth_img_1, depth_img_2], axis=1)\n",
    "\n",
    "    import cv2\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "    font_scale = 1.0\n",
    "    color_text = (255, 255, 255)\n",
    "    thickness = 2\n",
    "\n",
    "    # Original\n",
    "    cv2.putText(concat_img, \"Original\", (10, 30), font, font_scale, color_text, thickness, cv2.LINE_AA)\n",
    "    # label1\n",
    "    label_x = w + 10\n",
    "    cv2.putText(concat_img, label1, (label_x, 30), font, font_scale, color_text, thickness, cv2.LINE_AA)\n",
    "    # label2\n",
    "    label_x2 = w*2 + 10\n",
    "    cv2.putText(concat_img, label2, (label_x2, 30), font, font_scale, color_text, thickness, cv2.LINE_AA)\n",
    "\n",
    "    return concat_img\n",
    "\n",
    "def create_difference_map(depth_v2_color, distill_color):\n",
    "    \"\"\"\n",
    "    여기서는 Distill 결과가 단순 컬러 이미지이므로,\n",
    "    v2_color vs distill_color 간 RGB 차이만 계산\n",
    "    \"\"\"\n",
    "    diff_img = cv2.absdiff(depth_v2_color, distill_color)\n",
    "    return diff_img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [2] run_depth_comparison_on_input_folder (수정됨)\n",
    "def run_depth_comparison_on_input_folder():\n",
    "    \"\"\"\n",
    "    1) Distill-Any-Depth(data/input 폴더) 추론\n",
    "    2) Depth-Anything-V2 모델 로드 -> data/input 폴더 이미지 각각 inference\n",
    "    3) 비교 결과 /compare_depth 폴더에 저장 (3-way, diff)\n",
    "    * 저장 시 절대경로를 print\n",
    "    \"\"\"\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    # 1) Distill-Any-Depth 전체 추론 (결과 이미지도 compare_depth에 distill_ 접두어로 복사)\n",
    "    run_distill_any_depth()\n",
    "\n",
    "    # 2) Depth-Anything-V2 로드\n",
    "    model_v2 = load_depth_anything_v2(checkpoint_path=DEPTH_ANYTHING_CHECKPOINT_VITL, device=device)\n",
    "\n",
    "    # compare_depth 폴더 준비\n",
    "    os.makedirs(COMPARE_DEPTH_SAVE_DIR, exist_ok=True)\n",
    "\n",
    "    # data/input 폴더 내 이미지\n",
    "    input_files = [f for f in os.listdir(INPUT_FOLDER)\n",
    "                   if f.lower().endswith(('.jpg','.jpeg','.png','.bmp','.webp'))]\n",
    "    input_files.sort()\n",
    "\n",
    "    full_paths = [os.path.join(INPUT_FOLDER, x) for x in input_files]\n",
    "\n",
    "    # 3) Depth-Anything-V2 inference\n",
    "    depth_results_v2 = run_depth_anything_on_images(model_v2, full_paths, device=device, resize_size=700)\n",
    "\n",
    "    # Distill 결과(기본적으로 output/large/image_logs/ 아래 da_sota_*.jpg)\n",
    "    distill_image_logs = os.path.join(DISTILL_OUTPUT_DIR, \"image_logs\")\n",
    "    if not os.path.exists(distill_image_logs):\n",
    "        print(f\"[경고] {distill_image_logs} 폴더가 없어 Distill 결과를 찾을 수 없습니다.\")\n",
    "        return\n",
    "\n",
    "    distill_outputs = [f for f in os.listdir(distill_image_logs) if f.startswith(\"da_sota_\")]\n",
    "    distill_outputs.sort()\n",
    "\n",
    "    for idx, in_file in enumerate(input_files):\n",
    "        in_path = os.path.join(INPUT_FOLDER, in_file)\n",
    "        origin_pil = Image.open(in_path).convert(\"RGB\")\n",
    "        origin_np = np.array(origin_pil)\n",
    "        h, w, c = origin_np.shape\n",
    "\n",
    "        # Depth-Anything-V2 depth map -> 컬러맵\n",
    "        depth_map_v2 = depth_results_v2.get(in_file, None)\n",
    "        if depth_map_v2 is None:\n",
    "            print(f\"[오류] V2 depth 결과 없음: {in_file}\")\n",
    "            continue\n",
    "\n",
    "        depth_color_v2 = depth_to_colormap(depth_map_v2, cmap_name=\"plasma\")\n",
    "        depth_color_v2 = cv2.resize(depth_color_v2, (w, h), interpolation=cv2.INTER_LINEAR)\n",
    "\n",
    "        # Distill color (기존 로직, 단순 매핑)\n",
    "        if idx < len(distill_outputs):\n",
    "            distill_file = distill_outputs[idx]\n",
    "        else:\n",
    "            print(f\"[오류] Distill 결과 인덱스 초과 idx={idx}\")\n",
    "            continue\n",
    "        distill_img_path = os.path.join(distill_image_logs, distill_file)\n",
    "        if not os.path.isfile(distill_img_path):\n",
    "            print(f\"[오류] Distill 파일 없음: {distill_img_path}\")\n",
    "            continue\n",
    "\n",
    "        distill_pil = Image.open(distill_img_path).convert(\"RGB\")\n",
    "        distill_np = np.array(distill_pil)\n",
    "        distill_np = cv2.resize(distill_np, (w, h), interpolation=cv2.INTER_LINEAR)\n",
    "\n",
    "        # 차이맵\n",
    "        diff_img = create_difference_map(depth_color_v2, distill_np)\n",
    "\n",
    "        # 3-way concat\n",
    "        compare_3way = create_comparison_image(origin_np, depth_color_v2, distill_np,\n",
    "                                               \"DepthAnyV2\", \"Distill\")\n",
    "\n",
    "        out_name_base = os.path.splitext(in_file)[0]\n",
    "        out_v2_path = os.path.join(COMPARE_DEPTH_SAVE_DIR, f\"{out_name_base}_v2.png\")\n",
    "        out_distill_path = os.path.join(COMPARE_DEPTH_SAVE_DIR, f\"{out_name_base}_distill.png\")\n",
    "        out_compare_path = os.path.join(COMPARE_DEPTH_SAVE_DIR, f\"{out_name_base}_compare.png\")\n",
    "        out_diff_path = os.path.join(COMPARE_DEPTH_SAVE_DIR, f\"{out_name_base}_diff.png\")\n",
    "\n",
    "        cv2.imwrite(out_v2_path, cv2.cvtColor(depth_color_v2, cv2.COLOR_RGB2BGR))\n",
    "        cv2.imwrite(out_distill_path, cv2.cvtColor(distill_np, cv2.COLOR_RGB2BGR))\n",
    "        cv2.imwrite(out_compare_path, cv2.cvtColor(compare_3way, cv2.COLOR_RGB2BGR))\n",
    "        cv2.imwrite(out_diff_path, cv2.cvtColor(diff_img, cv2.COLOR_RGB2BGR))\n",
    "\n",
    "        # 절대경로로 표시\n",
    "        print(f\"[{in_file}] => 비교 결과 저장 완료\")\n",
    "        print(\"  -> DepthAny-V2:\", os.path.abspath(out_v2_path))\n",
    "        print(\"  -> Distill color:\", os.path.abspath(out_distill_path))\n",
    "        print(\"  -> Compare 3-way:\", os.path.abspath(out_compare_path))\n",
    "        print(\"  -> Diff map:\", os.path.abspath(out_diff_path))\n",
    "\n",
    "    print(\"=== 모든 이미지 비교 작업을 마쳤습니다! ===\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------\n",
    "# 7) 실행 (주피터 노트북에서 이 셀을 실행하거나, python 스크립트로도 가능)\n",
    "# -------------------------------------------------------------\n",
    "if __name__ == \"__main__\":\n",
    "    run_depth_comparison_on_input_folder()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sam2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}