{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 1. 초기 설정 및 라이브러리 임포트\n",
    "이 셀에서는 필요한 라이브러리를 임포트하고, 실행 환경(device) 및\n",
    "결과물(이미지, npy 등)을 저장할 폴더를 자동으로 생성합니다.\n",
    "출력 파일들은 절대경로(예: /ltb/media/ltb/90887173887158A4/Users/ltb/sam2_output/exp, exp2, …)에 저장됩니다.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "import os\n",
    "import cv2\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image, ImageDraw\n",
    "from datetime import datetime\n",
    "import json\n",
    "import re\n",
    "import time\n",
    "\n",
    "# GPU가 있다면 GPU 사용, 없으면 CPU 사용\n",
    "def select_device():\n",
    "    if torch.cuda.is_available():\n",
    "        return torch.device(\"cuda\")\n",
    "    else:\n",
    "        return torch.device(\"cpu\")\n",
    "\n",
    "device = select_device()\n",
    "print(f\"[초기 설정] 사용 중인 device: {device}\")\n",
    "\n",
    "# 결과 저장 폴더 자동 생성 (덮어쓰지 않도록 exp, exp2, ... 형식)\n",
    "def get_new_experiment_folder(base_path=\"/ltb/media/ltb/90887173887158A4/Users/ltb/sam2_output\", prefix=\"exp\"):\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "    idx = 1\n",
    "    while True:\n",
    "        folder_name = prefix if idx == 1 else f\"{prefix}{idx}\"\n",
    "        exp_dir = os.path.join(base_path, folder_name)\n",
    "        if not os.path.exists(exp_dir):\n",
    "            os.makedirs(exp_dir)\n",
    "            return exp_dir\n",
    "        idx += 1\n",
    "\n",
    "EXPERIMENT_DIR = get_new_experiment_folder()\n",
    "print(f\"[초기 설정] 실험 결과 폴더: {EXPERIMENT_DIR}\")\n",
    "\n",
    "# 현재 날짜 문자열 (YYYYmmdd), 파일명 등에 사용\n",
    "DATE_STR = datetime.now().strftime(\"%Y%m%d\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 2. 데이터 불러오기 및 첫 프레임 추출\n",
    "입력 영상(절대 경로 지정)의 첫 프레임을 추출하여 이미지로 저장합니다.\n",
    "또한 영상 해상도(가로 × 세로)를 출력하여, 후속 포인트 입력 전에 확인할 수 있도록 합니다.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "# 입력 영상 파일 경로 – 환경에 맞게 수정하세요.\n",
    "video_path = \"/ltb/media/ltb/90887173887158A4/Users/ltb/주피터노트북/wire_real_og_40min_split.mp4\"\n",
    "print(f\"[데이터 불러오기] 입력 영상: {video_path}\")\n",
    "\n",
    "# cv2로 영상 열기\n",
    "cap = cv2.VideoCapture(video_path)\n",
    "if not cap.isOpened():\n",
    "    raise IOError(f\"영상 열기 실패: {video_path}\")\n",
    "\n",
    "# 영상 해상도 출력\n",
    "width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "print(f\"[데이터 불러오기] 영상 해상도: {width} x {height}\")\n",
    "\n",
    "# 첫 프레임 추출\n",
    "ret, first_frame = cap.read()\n",
    "cap.release()\n",
    "if not ret:\n",
    "    raise ValueError(\"첫 프레임 추출 실패\")\n",
    "\n",
    "# BGR -> RGB 변환 후 PIL 이미지 변환\n",
    "first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)\n",
    "first_frame_img = Image.fromarray(first_frame_rgb)\n",
    "\n",
    "# 첫 프레임 저장 (절대경로에 저장)\n",
    "first_frame_save_path = os.path.join(EXPERIMENT_DIR, f\"{DATE_STR}_first_frame.png\")\n",
    "first_frame_img.save(first_frame_save_path)\n",
    "print(f\"[데이터 불러오기] 첫 프레임 저장 완료: {first_frame_save_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 3. 포인트 입력 및 시각화\n",
    "사용자가 수동으로 지정한 포인트(positive, negative)를 입력하고,\n",
    "첫 프레임 이미지에 오버레이하여 시각적으로 확인합니다.\n",
    "\n",
    "- **positive_points**: 세그멘테이션 대상 영역을 지정하는 좌표 (예: [(x, y), ...])\n",
    "- **negative_points**: 제외할 영역 (필요 시; 없으면 빈 리스트)\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "# 예시 좌표 – 필요한 경우 사용자가 직접 수정하세요.\n",
    "positive_points = [(300, 100), (300, 200), (300, 350)]\n",
    "negative_points = []  # 예: [(800, 1430)]\n",
    "\n",
    "print(f\"[포인트 입력] positive_points: {positive_points}\")\n",
    "print(f\"[포인트 입력] negative_points: {negative_points}\")\n",
    "\n",
    "# 첫 프레임 이미지에 포인트 오버레이 (PIL 사용)\n",
    "vis_img = first_frame_img.convert(\"RGB\")\n",
    "draw = ImageDraw.Draw(vis_img)\n",
    "\n",
    "# positive 포인트: 파란색 원과 라벨 표시\n",
    "for idx, pt in enumerate(positive_points):\n",
    "    x, y = pt\n",
    "    r = 8\n",
    "    draw.ellipse((x - r, y - r, x + r, y + r), outline=\"blue\", width=2)\n",
    "    draw.text((x + r, y - r), f\"[+{idx}]\", fill=\"blue\")\n",
    "\n",
    "# negative 포인트: 빨간색 십자(×)와 라벨 표시\n",
    "for idx, pt in enumerate(negative_points):\n",
    "    x, y = pt\n",
    "    r = 8\n",
    "    draw.line((x - r, y - r, x + r, y + r), fill=\"red\", width=2)\n",
    "    draw.line((x - r, y + r, x + r, y - r), fill=\"red\", width=2)\n",
    "    draw.text((x + r, y - r), f\"[-{idx}]\", fill=\"red\")\n",
    "\n",
    "# 결과 이미지 저장 및 인라인 시각화\n",
    "vis_img_path = os.path.join(EXPERIMENT_DIR, f\"{DATE_STR}_first_frame_with_points.png\")\n",
    "vis_img.save(vis_img_path)\n",
    "print(f\"[포인트 시각화] 포인트 오버레이 이미지 저장: {vis_img_path}\")\n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.imshow(vis_img)\n",
    "plt.title(\"First Frame with Prompt Points\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 4. 세그멘테이션 모드 선택\n",
    "이 셀에서는 포인트 기반 세그멘테이션(\"point\")과\n",
    "박스 기반 세그멘테이션(\"box\") 중 하나를 선택할 수 있습니다.\n",
    "- **SEGMENTATION_MODE = \"point\"** → 사용자가 입력한 좌표를 그대로 사용합니다.\n",
    "- **SEGMENTATION_MODE = \"box\"** → positive 포인트들을 바운딩박스로 변환하여 사용합니다.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "# \"point\" 또는 \"box\" 중 원하는 모드로 설정하세요.\n",
    "SEGMENTATION_MODE = \"point\"  # 예: \"point\"\n",
    "# SEGMENTATION_MODE = \"box\"  # 예: \"box\"\n",
    "\n",
    "print(f\"[세그멘테이션 모드] 현재 모드: {SEGMENTATION_MODE}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 5. SAM2 모델 로드\n",
    "이 셀에서는 SAM2의 포인트(또는 박스) 프롬프트 기능을 사용하기 위해 모델을 로드합니다.\n",
    "현재 노트북 위치는:\n",
    "    /ltb/media/ltb/90887173887158A4/Users/ltb/주피터노트북/sam2.1로_마스크추출_간소화.ipynb\n",
    "\n",
    "예시 경로:\n",
    "- SAM2 config 파일: \"../회사용/4090_옮길거/workspace/sam2_repo/sam2/configs/sam2.1/sam2.1_hiera_l.yaml\"\n",
    "- SAM2 checkpoint: \"../회사용/4090_옮길거/workspace/sam2_repo/checkpoints/sam2.1_hiera_large.pt\"\n",
    "\n",
    "**중요:** config 파일 경로 앞에 \"file://\"를 붙여 절대경로임을 Hydra에 명시합니다.\n",
    "경로는 환경에 맞게 수정하세요.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "# 현재 노트북의 상위 폴더를 기준으로 경로 설정 (상대경로 → 절대경로)\n",
    "BASE_DIR = os.path.abspath(os.path.join(os.getcwd(), \"..\"))\n",
    "SAM2_CONFIG_REL = \"회사용/4090_옮길거/workspace/sam2_repo/sam2/configs/sam2.1/sam2.1_hiera_l.yaml\"\n",
    "SAM2_CHECKPOINT_REL = \"회사용/4090_옮길거/workspace/sam2_repo/checkpoints/sam2.1_hiera_large.pt\"\n",
    "\n",
    "# \"file://\" 접두사 추가하여 config 파일 절대경로 생성\n",
    "SAM2_CONFIG_FILE = \"file://\" + os.path.join(BASE_DIR, SAM2_CONFIG_REL)\n",
    "SAM2_CHECKPOINT_PATH = os.path.join(BASE_DIR, SAM2_CHECKPOINT_REL)\n",
    "\n",
    "print(f\"[SAM2 설정] Config 파일: {SAM2_CONFIG_FILE}\")\n",
    "print(f\"[SAM2 설정] Checkpoint 파일: {SAM2_CHECKPOINT_PATH}\")\n",
    "\n",
    "# SAM2 모델 로드를 위한 함수 (SAM2 리포지토리 내 build_sam 모듈 사용)\n",
    "try:\n",
    "    from sam2.build_sam import build_sam2_video_predictor\n",
    "except ImportError:\n",
    "    raise ImportError(\"SAM2 모듈을 찾을 수 없습니다. SAM2 코드를 설치하고 PYTHONPATH에 추가하세요.\")\n",
    "\n",
    "# SAM2 모델 로드 (eval 모드)\n",
    "predictor = build_sam2_video_predictor(\n",
    "    config_file=SAM2_CONFIG_FILE,\n",
    "    ckpt_path=SAM2_CHECKPOINT_PATH,\n",
    "    device=device,\n",
    "    mode=\"eval\",\n",
    "    apply_postprocessing=True\n",
    ")\n",
    "print(\"[SAM2 모델] 로드 완료\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 6. 세그멘테이션 추론 및 결과 저장\n",
    "이 셀에서는 첫 프레임 이미지와 사용자가 지정한 포인트(또는 박스)를 SAM2 모델에 입력하여,\n",
    "세그멘테이션 마스크를 추론합니다.\n",
    "추론된 마스크는 npy 파일로 저장되고, 원본 이미지에 투명 파란색 오버레이한 결과도 저장 및 인라인 시각화됩니다.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "# 첫 프레임 이미지를 numpy 배열(RGB)로 변환\n",
    "first_frame_np = np.array(first_frame_img.convert(\"RGB\"))\n",
    "h, w = first_frame_np.shape[:2]\n",
    "\n",
    "if SEGMENTATION_MODE == \"point\":\n",
    "    # 포인트 기반: 좌표를 float32로 변환 후 이미지 크기 기준 0~1로 정규화\n",
    "    positive_points_np = np.array(positive_points, dtype=np.float32) / np.array([w, h], dtype=np.float32)\n",
    "    negative_points_np = (np.array(negative_points, dtype=np.float32) / np.array([w, h], dtype=np.float32)\n",
    "                          if negative_points else np.empty((0, 2), dtype=np.float32))\n",
    "    \n",
    "    positive_labels = [1] * len(positive_points)\n",
    "    negative_labels = [0] * len(negative_points)\n",
    "    all_points = (np.concatenate([positive_points_np, negative_points_np], axis=0)\n",
    "                  if len(negative_points_np) > 0 else positive_points_np)\n",
    "    all_labels = positive_labels + negative_labels\n",
    "\n",
    "    print(\"[세그멘테이션 - 포인트] 정규화된 포인트 (좌표 0~1):\", all_points)\n",
    "    print(\"[세그멘테이션 - 포인트] 포인트 라벨:\", all_labels)\n",
    "    \n",
    "    try:\n",
    "        masks, scores, logits = predictor.predict(\n",
    "            image=first_frame_np,\n",
    "            point_coords=all_points,\n",
    "            point_labels=np.array(all_labels),\n",
    "            multimask_output=False\n",
    "        )\n",
    "    except Exception as e:\n",
    "        raise RuntimeError(f\"SAM2 예측 중 오류 발생 (포인트): {e}\")\n",
    "    \n",
    "elif SEGMENTATION_MODE == \"box\":\n",
    "    # 박스 기반: positive 포인트로부터 바운딩박스 생성 (픽셀 단위)\n",
    "    positive_points_arr = np.array(positive_points, dtype=np.float32)\n",
    "    min_xy = positive_points_arr.min(axis=0)\n",
    "    max_xy = positive_points_arr.max(axis=0)\n",
    "    bbox = [min_xy[0], min_xy[1], max_xy[0], max_xy[1]]\n",
    "    bbox_normalized = np.array(bbox, dtype=np.float32) / np.array([w, h, w, h], dtype=np.float32)\n",
    "    print(f\"[세그멘테이션 - 박스] 생성된 바운딩박스 (픽셀): {bbox}\")\n",
    "    print(f\"[세그멘테이션 - 박스] 정규화된 바운딩박스: {bbox_normalized}\")\n",
    "    \n",
    "    try:\n",
    "        masks, scores, logits = predictor.predict(\n",
    "            image=first_frame_np,\n",
    "            point_coords=None,\n",
    "            point_labels=None,\n",
    "            box=np.expand_dims(bbox_normalized, axis=0),\n",
    "            multimask_output=False\n",
    "        )\n",
    "    except Exception as e:\n",
    "        raise RuntimeError(f\"SAM2 예측 중 오류 발생 (박스): {e}\")\n",
    "else:\n",
    "    raise ValueError(\"SEGMENTATION_MODE 값이 잘못되었습니다. 'point' 또는 'box'로 설정하세요.\")\n",
    "\n",
    "if masks is None or len(masks) == 0:\n",
    "    raise ValueError(\"세그멘테이션 마스크가 생성되지 않았습니다.\")\n",
    "mask_result = masks[0]\n",
    "print(f\"[세그멘테이션] 추론 완료, 마스크 shape: {mask_result.shape}\")\n",
    "\n",
    "# 마스크 결과 저장 (npy 파일)\n",
    "mask_save_path = os.path.join(EXPERIMENT_DIR, f\"{DATE_STR}_mask_result.npy\")\n",
    "np.save(mask_save_path, mask_result)\n",
    "print(f\"[세그멘테이션] 마스크 npy 저장 완료: {mask_save_path}\")\n",
    "\n",
    "# 원본 이미지에 투명 파란색 오버레이 (시각화)\n",
    "overlay_color = (0, 0, 255)  # RGB\n",
    "alpha = 0.4\n",
    "\n",
    "overlay_img = first_frame_img.convert(\"RGBA\")\n",
    "mask_img = Image.fromarray((mask_result * 255).astype(np.uint8)).convert(\"L\")\n",
    "color_img = Image.new(\"RGBA\", (w, h), overlay_color + (0,))\n",
    "mask_np = np.array(mask_img)\n",
    "alpha_channel = np.where(mask_np > 127, int(255 * alpha), 0).astype(np.uint8)\n",
    "color_np = np.array(color_img)\n",
    "color_np[..., 3] = alpha_channel\n",
    "color_img = Image.fromarray(color_np, mode=\"RGBA\")\n",
    "\n",
    "combined_img = Image.alpha_composite(overlay_img, color_img)\n",
    "\n",
    "combined_img_save_path = os.path.join(EXPERIMENT_DIR, f\"{DATE_STR}_first_frame_mask_overlay.png\")\n",
    "combined_img.save(combined_img_save_path)\n",
    "print(f\"[세그멘테이션] 오버레이 이미지 저장 완료: {combined_img_save_path}\")\n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.imshow(combined_img)\n",
    "plt.title(\"Segmentation Result Overlay\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% [markdown]\n",
    "\"\"\"\n",
    "# 7. 최종 결과 확인 및 로그 출력\n",
    "세그멘테이션 마스크와 오버레이 이미지의 저장 경로, 그리고 마스크의 통계 정보를 출력합니다.\n",
    "이 정보를 통해 결과 파일들을 확인하고, 후속 분석이나 디버깅에 활용할 수 있습니다.\n",
    "\"\"\"\n",
    "\n",
    "# %% [code]\n",
    "print(\"=== 최종 결과 파일 ===\")\n",
    "print(f\"1) 첫 프레임 이미지: {first_frame_save_path}\")\n",
    "print(f\"2) 포인트 오버레이 이미지: {vis_img_path}\")\n",
    "print(f\"3) 세그멘테이션 마스크 (npy): {mask_save_path}\")\n",
    "print(f\"4) 마스크 오버레이 이미지: {combined_img_save_path}\")\n",
    "\n",
    "mask_loaded = np.load(mask_save_path)\n",
    "print(f\"마스크 min: {mask_loaded.min()}, max: {mask_loaded.max()}, mean: {mask_loaded.mean():.4f}\")\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
}