{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [0]\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",
    "# 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 / V2 결과 비교를 저장할 상위 폴더\n",
    "COMPARE_DEPTH_BASE = \"/ltb/media/ltb/90887173887158A4/Users/ltb/compare_depth\"\n",
    "\n",
    "# Distill 입력 이미지 폴더(동일)\n",
    "INPUT_FOLDER = DISTILL_INPUT_DIR\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [1]\n",
    "def get_new_experiment_folder(base_path, prefix=\"exp\"):\n",
    "    \"\"\"\n",
    "    base_path 아래에 prefix로 시작하는 폴더 생성.\n",
    "    이미 exp가 있으면 exp2, exp3... 식으로 증가.\n",
    "    \"\"\"\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "    i = 0\n",
    "    while True:\n",
    "        i += 1\n",
    "        if i==1:\n",
    "            folder_name = prefix\n",
    "        else:\n",
    "            folder_name = f\"{prefix}{i}\"\n",
    "        new_dir = os.path.join(base_path, folder_name)\n",
    "        if not os.path.exists(new_dir):\n",
    "            os.makedirs(new_dir)\n",
    "            return new_dir\n",
    "    # 반환 후, ex) /ltb/media/ltb/compare_depth/exp2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [2]\n",
    "def run_distill_any_depth():\n",
    "    \"\"\"\n",
    "    1) Distill-Any-Depth 추론\n",
    "    2) /output/large/image_logs -> compare_depth 폴더로 복사\n",
    "    3) 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",
    "            path2 = os.path.join(DISTILL_OUTPUT_DIR, old_file)\n",
    "            if os.path.isfile(path2):\n",
    "                os.remove(path2)\n",
    "            else:\n",
    "                shutil.rmtree(path2, ignore_errors=True)\n",
    "    else:\n",
    "        os.makedirs(DISTILL_OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "    # 실행\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [3]\n",
    "def load_depth_anything_v2(checkpoint_path, device=\"cuda\"):\n",
    "    \"\"\"\n",
    "    V2 모델 로드 (vitl)\n",
    "    \"\"\"\n",
    "    if DEPTH_ANYTHING_DIR not in sys.path:\n",
    "        sys.path.append(DEPTH_ANYTHING_DIR)\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.to(device)\n",
    "    model.eval()\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [4]\n",
    "def run_depth_anything_on_images(model, input_images, device=\"cuda\", resize_size=700):\n",
    "    \"\"\"\n",
    "    input_images 리스트 각 이미지에 대해 Depth-Anything-V2 inference\n",
    "    return: {파일명: depth_map}\n",
    "    \"\"\"\n",
    "    depth_results = {}\n",
    "    with torch.no_grad():\n",
    "        for path in input_images:\n",
    "            nm = os.path.basename(path)\n",
    "            pil_img = Image.open(path).convert('RGB')\n",
    "            arr = np.array(pil_img)\n",
    "            dmap = model.infer_image(arr, resize_size)\n",
    "            depth_results[nm] = dmap.squeeze()\n",
    "    return depth_results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [5]\n",
    "def normalize_minmax(depth_map):\n",
    "    mn, mx = depth_map.min(), depth_map.max()\n",
    "    if mx - mn < 1e-6:\n",
    "        return np.zeros_like(depth_map, dtype=np.float32)\n",
    "    return (depth_map - mn)/(mx - mn)\n",
    "\n",
    "def depth_to_colormap(depth_map, cmap_name=\"plasma\"):\n",
    "    \"\"\"\n",
    "    depth_map: 2D numpy\n",
    "    return: 8bit RGB (H,W,3)\n",
    "    \"\"\"\n",
    "    d_norm = normalize_minmax(depth_map)\n",
    "    cmap = plt.get_cmap(cmap_name)\n",
    "    colored = cmap(d_norm)[:,:,:3]  # RGBA => RGB\n",
    "    return (colored*255).astype(np.uint8)\n",
    "\n",
    "def create_comparison_image(orig, c1, c2, label1=\"ModelA\", label2=\"ModelB\"):\n",
    "    \"\"\"\n",
    "    orig, c1, c2 shape: (H,W,3)\n",
    "    \"\"\"\n",
    "    h,w,_=orig.shape\n",
    "    merged = np.concatenate([orig, c1, c2], axis=1)\n",
    "    import cv2\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "    cv2.putText(merged, \"Original\", (10,30), font,1.0,(255,255,255),2,cv2.LINE_AA)\n",
    "    cv2.putText(merged, label1, (w+10,30), font,1.0,(255,255,255),2,cv2.LINE_AA)\n",
    "    cv2.putText(merged, label2, (w*2+10,30), font,1.0,(255,255,255),2,cv2.LINE_AA)\n",
    "    return merged\n",
    "\n",
    "def create_difference_map(img1, img2):\n",
    "    \"\"\"\n",
    "    단순 RGB diff\n",
    "    \"\"\"\n",
    "    return cv2.absdiff(img1, img2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [6]\n",
    "def get_new_experiment_folder(base_path, prefix=\"exp\"):\n",
    "    \"\"\"\n",
    "    base_path 아래 prefix폴더 생성. 이미 있으면 exp2, exp3...\n",
    "    \"\"\"\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "    idx=0\n",
    "    while True:\n",
    "        idx+=1\n",
    "        if idx==1:\n",
    "            folder_name=prefix\n",
    "        else:\n",
    "            folder_name=f\"{prefix}{idx}\"\n",
    "        new_dir=os.path.join(base_path, folder_name)\n",
    "        if not os.path.exists(new_dir):\n",
    "            os.makedirs(new_dir)\n",
    "            return new_dir\n",
    "    # return ...\n",
    "\n",
    "def prepare_exp_subdirs(exp_dir):\n",
    "    \"\"\"\n",
    "    exp_dir 아래에 v2, distill, compare, diff 4개 하위폴더 생성\n",
    "    return dict of subdir paths\n",
    "    \"\"\"\n",
    "    subdirs={}\n",
    "    for name in [\"v2\",\"distill\",\"compare\",\"diff\"]:\n",
    "        sd = os.path.join(exp_dir, name)\n",
    "        os.makedirs(sd, exist_ok=True)\n",
    "        subdirs[name]=sd\n",
    "    return subdirs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [7]\n",
    "def run_depth_comparison_on_input_folder():\n",
    "    \"\"\"\n",
    "    1) Distill 추론 => /output/large/image_logs/*.jpg 생성\n",
    "    2) exp_dir 만들고 distill logs도 exp_dir/distill 에 복사\n",
    "    3) Depth-Anything-V2 inference => exp_dir/v2\n",
    "    4) compare => exp_dir/compare, diff => exp_dir/diff\n",
    "    \"\"\"\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    # 1) Distill-Any-Depth 추론\n",
    "    run_distill_any_depth()\n",
    "\n",
    "    # 2) 새 exp 폴더 만들기\n",
    "    exp_dir = get_new_experiment_folder(COMPARE_DEPTH_BASE, prefix=\"exp\")\n",
    "    subdirs=prepare_exp_subdirs(exp_dir)\n",
    "    print(f\"[INFO] New experiment folder => {exp_dir}\")\n",
    "\n",
    "    # 2-1) Distill / output/large/image_logs => exp_dir/distill 복사\n",
    "    distill_logs = os.path.join(DISTILL_OUTPUT_DIR, \"image_logs\")\n",
    "    if not os.path.exists(distill_logs):\n",
    "        print(\"[경고] distill image_logs 없음.\")\n",
    "    else:\n",
    "        files = sorted([f for f in os.listdir(distill_logs) if f.startswith(\"da_sota_\")])\n",
    "        for f in files:\n",
    "            srcp=os.path.join(distill_logs, f)\n",
    "            dstp=os.path.join(subdirs[\"distill\"], f)\n",
    "            shutil.copy2(srcp, dstp)\n",
    "        print(f\"[INFO] Distill logs => {subdirs['distill']}\")\n",
    "\n",
    "    # 3) Depth-Anything-V2 로드 & inference\n",
    "    model_v2 = load_depth_anything_v2(DEPTH_ANYTHING_CHECKPOINT_VITL, device=device)\n",
    "\n",
    "    # 입력 파일들\n",
    "    input_files = sorted([f for f in os.listdir(INPUT_FOLDER)\n",
    "                          if f.lower().endswith(('.jpg','.jpeg','.png','.bmp','.webp'))])\n",
    "    full_paths = [os.path.join(INPUT_FOLDER, x) for x in input_files]\n",
    "    # V2 inference => dict( filename => depth_map )\n",
    "    depth_results_v2 = {}\n",
    "    with torch.no_grad():\n",
    "        for path in full_paths:\n",
    "            nm = os.path.basename(path)\n",
    "            pil_img = Image.open(path).convert('RGB')\n",
    "            arr = np.array(pil_img)\n",
    "            dmap = model_v2.infer_image(arr, 700)\n",
    "            depth_results_v2[nm]=dmap.squeeze()\n",
    "\n",
    "    # 3-1) v2 컬러맵 => exp_dir/v2\n",
    "    for nm, dmap in depth_results_v2.items():\n",
    "        h,w=dmap.shape\n",
    "        dcolor=depth_to_colormap(dmap, \"plasma\")\n",
    "        # 원본 해상도 => actually we know arr is h,w\n",
    "        outpath=os.path.join(subdirs[\"v2\"], nm.replace(\".\",\"_v2.\"))\n",
    "        cv2.imwrite(outpath, cv2.cvtColor(dcolor, cv2.COLOR_RGB2BGR))\n",
    "    print(f\"[INFO] Depth-Anything-V2 => {subdirs['v2']}\")\n",
    "\n",
    "    # 4) compare => exp_dir/compare, diff => exp_dir/diff\n",
    "    # Distill => da_sota_{idx}.jpg\n",
    "    # We'll match index-based approach: for i, in_file in enumerate(input_files) => distill => da_sota_{i}.jpg\n",
    "    for i, in_file in enumerate(input_files):\n",
    "        # e.g. 0 => da_sota_0.jpg\n",
    "        distill_jpg = f\"da_sota_{i}.jpg\"\n",
    "        # check if distill_jpg exists\n",
    "        distill_jpg_path = os.path.join(subdirs[\"distill\"], distill_jpg)\n",
    "        if not os.path.isfile(distill_jpg_path):\n",
    "            print(f\"[경고] Distill {distill_jpg} 없음 => 스킵\")\n",
    "            continue\n",
    "        # load distill, v2\n",
    "        nm = os.path.basename(in_file)\n",
    "        origin_pil=Image.open(os.path.join(INPUT_FOLDER, in_file)).convert(\"RGB\")\n",
    "        origin_np=np.array(origin_pil)\n",
    "        d_h,d_w,_=origin_np.shape\n",
    "\n",
    "        # v2 color\n",
    "        if nm not in depth_results_v2:\n",
    "            print(f\"[경고] V2 result 없음 => {nm}\")\n",
    "            continue\n",
    "        v2_dmap=depth_results_v2[nm]\n",
    "        # colormap\n",
    "        v2_color=depth_to_colormap(v2_dmap,\"plasma\")\n",
    "        v2_color=cv2.resize(v2_color,(d_w,d_h),cv2.INTER_LINEAR)\n",
    "\n",
    "        distill_img=cv2.imread(distill_jpg_path)\n",
    "        if distill_img is None:\n",
    "            print(f\"[경고] distill_img 로드 실패 => {distill_jpg_path}\")\n",
    "            continue\n",
    "        distill_img=cv2.cvtColor(distill_img, cv2.COLOR_BGR2RGB)\n",
    "        distill_img=cv2.resize(distill_img,(d_w,d_h),cv2.INTER_LINEAR)\n",
    "\n",
    "        # diff\n",
    "        diff_img = create_difference_map(v2_color, distill_img)\n",
    "        # compare 3way\n",
    "        compare_img = create_comparison_image(origin_np,v2_color,distill_img,\"V2\",\"Distill\")\n",
    "\n",
    "        # 저장\n",
    "        base_noext=os.path.splitext(in_file)[0]\n",
    "        # ex) origin= \"00008.jpg\" => base_noext=\"00008\"\n",
    "        out_compare=os.path.join(subdirs[\"compare\"], base_noext+\"_compare.png\")\n",
    "        out_diff   =os.path.join(subdirs[\"diff\"],    base_noext+\"_diff.png\")\n",
    "\n",
    "        cv2.imwrite(out_compare, cv2.cvtColor(compare_img, cv2.COLOR_RGB2BGR))\n",
    "        cv2.imwrite(out_diff,    cv2.cvtColor(diff_img,    cv2.COLOR_RGB2BGR))\n",
    "\n",
    "        print(f\"[{in_file}] => compare={out_compare}, diff={out_diff}\")\n",
    "\n",
    "    print(\"[INFO] 모든 이미지 비교 완료.\")\n",
    "    print(f\"[INFO] 결과 폴더 => {exp_dir}\")\n",
    "\n",
    "\n",
    "# %% [8] 메인\n",
    "if __name__==\"__main__\":\n",
    "    run_depth_comparison_on_input_folder()\n"
   ]
  }
 ],
 "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
}