{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":117682,"databundleVersionId":14443416,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport glob\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn.functional as F\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\n\n# ==========================================\n# 1. CONSTANTS & CONFIGURATION\n# ==========================================\n\n# Input Paths\nPATH = '/kaggle/input/vesuvius-challenge-surface-detection'\nOUTPUT_DIR = './processed_dataset_v1'\n\n# Data Processing Params\nIMG_SIZE = 128         # Final output size (Resize target)\nHALF_WINDOW = 10       # How many slices +/- from center to use\nNEIGHBORS = [-1, 0, 1] # 3-slice stack for 2.5D input\n\n# Device\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Load CSV to map Volume IDs to Scroll IDs\n# This is crucial for organizing output into scroll folders\ndf_train = pd.read_csv(f\"{PATH}/train.csv\")\nID_TO_SCROLL = dict(zip(df_train['id'], df_train['scroll_id']))\n\nprint(f\"Config: {DEVICE} | Output: {OUTPUT_DIR}\")\nprint(f\"Mapping: {len(ID_TO_SCROLL)} volumes mapped to scrolls.\")\n\n# ==========================================\n# 2. HELPER FUNCTIONS\n# ==========================================\n\ndef load_subvolume_indices(n_frames, mode='train'):\n    \"\"\"\n    Determines which slice indices we need to process.\n    Train/Valid: Only the middle chunk where the papyrus is likely located.\n    Test: All frames.\n    \"\"\"\n    if mode == 'train' or mode == 'valid':\n        mid_idx = n_frames // 2\n        start_idx = max(0, mid_idx - HALF_WINDOW)\n        end_idx = min(n_frames, mid_idx + HALF_WINDOW)\n        return range(start_idx, end_idx)\n    else:\n        return range(n_frames)\n\ndef load_tiff_slice_range(path, start_idx, end_idx, n_frames):\n    \"\"\"\n    Optimization: Loads only the required range of frames into RAM at once.\n    Returns: (Depth, H, W) array + the start offset.\n    \"\"\"\n    # Load slightly wider buffer for neighbors (-1 / +1)\n    load_start = max(0, start_idx - 1)\n    load_end = min(n_frames, end_idx + 1)\n    \n    imgs = []\n    with Image.open(path) as img:\n        for i in range(load_start, load_end):\n            img.seek(i)\n            imgs.append(np.array(img))\n            \n    stack = np.stack(imgs, axis=0) \n    return stack, load_start\n\ndef process_and_save_volume(img_id, mode='train'):\n    \"\"\"\n    Loads volume, looks up its Scroll ID, and saves to scroll-specific folder.\n    \"\"\"\n    folder = 'train_images' if mode != 'test' else 'test_images'\n    img_path = f\"{PATH}/{folder}/{img_id}.tif\"\n    \n    # 1. Determine Output Folder based on Scroll ID\n    if mode == 'train':\n        # Look up scroll ID from CSV map\n        try:\n            scroll_id = ID_TO_SCROLL[int(img_id)]\n        except KeyError:\n            print(f\"Warning: ID {img_id} not found in CSV. Using 'unknown'.\")\n            scroll_id = \"unknown\"\n            \n        # Structure: output/train/{scroll_id}/images/\n        save_dir_img = f\"{OUTPUT_DIR}/{mode}/{scroll_id}/images\"\n        save_dir_msk = f\"{OUTPUT_DIR}/{mode}/{scroll_id}/masks\"\n    else:\n        # Test set usually goes into a single folder\n        save_dir_img = f\"{OUTPUT_DIR}/{mode}/images\"\n    \n    # Create Directories\n    os.makedirs(save_dir_img, exist_ok=True)\n    if mode == 'train': os.makedirs(save_dir_msk, exist_ok=True)\n\n    # 2. Get Volume Metadata\n    with Image.open(img_path) as img:\n        n_frames = img.n_frames\n        \n    target_indices = load_subvolume_indices(n_frames, mode)\n    if len(target_indices) == 0: return 0\n\n    # 3. Load Image Data (Optimized)\n    start_target = target_indices[0]\n    end_target = target_indices[-1] + 1\n    vol_arr, buffer_start_offset = load_tiff_slice_range(img_path, start_target, end_target, n_frames)\n    \n    # 4. Load Mask Data (Train only)\n    mask_arr = None\n    if mode == 'train':\n        mask_path = f\"{PATH}/train_labels/{img_id}.tif\"\n        m_imgs = []\n        with Image.open(mask_path) as m:\n            for i in target_indices:\n                m.seek(i)\n                m_imgs.append(np.array(m))\n        mask_arr = np.stack(m_imgs, axis=0)\n\n    # 5. Iterate & Save Patches\n    count = 0\n    for i, slice_idx in enumerate(tqdm(target_indices, desc=f\"Vol {img_id}\", leave=False)):\n        \n        # --- A. Construct 3-Slice Stack ---\n        local_center = slice_idx - buffer_start_offset\n        slices = []\n        for offset in NEIGHBORS: # [-1, 0, 1]\n            nb_idx = np.clip(local_center + offset, 0, vol_arr.shape[0] - 1)\n            slices.append(vol_arr[nb_idx])\n            \n        stack = np.stack(slices, axis=0).astype(np.float32)\n\n        # --- B. Z-Score Normalization ---\n        stack = (stack - stack.mean()) / (stack.std() + 1e-6)\n        \n        # --- C. Resize (GPU) ---\n        tensor_img = torch.tensor(stack).to(DEVICE)\n        tensor_img = F.interpolate(\n            tensor_img.unsqueeze(0), \n            size=(IMG_SIZE, IMG_SIZE), \n            mode='bilinear', \n            align_corners=False\n        ).squeeze(0).cpu().numpy()\n\n        # Save Name: {volume_id}_{slice_index}.npy\n        save_id = f\"{img_id}_{slice_idx}\"\n        np.save(f\"{save_dir_img}/{save_id}.npy\", tensor_img)\n        \n        # --- D. Process Mask (Train Only) ---\n        if mode == 'train':\n            curr_mask = mask_arr[i]\n            mask_tensor = torch.tensor(curr_mask).long().to(DEVICE)\n            \n            # Resize Mask (Nearest Neighbor)\n            mask_tensor = F.interpolate(\n                mask_tensor.unsqueeze(0).unsqueeze(0).float(), \n                size=(IMG_SIZE, IMG_SIZE), \n                mode='nearest'\n            ).long().squeeze().cpu().numpy()\n            \n            np.save(f\"{save_dir_msk}/{save_id}.npy\", mask_tensor)\n            \n        count += 1\n    return count\n\n# ==========================================\n# 3. MAIN EXECUTION LOOP\n# ==========================================\n\n# Get list of Image IDs\ntrain_files = sorted(glob.glob(f\"{PATH}/train_images/*.tif\"))\ntrain_ids = [os.path.basename(f).split('.')[0] for f in train_files]\n\nprint(f\"Found Train IDs: {train_ids}\")\n\ntotal_saved = 0\n\nfor img_id in tqdm(train_ids, desc=\"Total Progress\"):\n    n_saved = process_and_save_volume(img_id, mode='train')\n    total_saved += n_saved\n\nprint(f\"\\nPreprocessing Complete. Saved {total_saved} samples to {OUTPUT_DIR}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T06:41:03.352863Z","iopub.execute_input":"2025-12-07T06:41:03.353221Z","execution_failed":"2025-12-07T06:41:27.041Z"}},"outputs":[],"execution_count":null}]}