{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.6.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n\n# Use the kagglehub client library to attach Kaggle resources like competitions, datasets, and models to your session\n# Learn more about kagglehub: https://github.com/Kaggle/kagglehub/blob/main/README.md\n\nimport kagglehub\n# kagglehub.dataset_download('<owner>/<dataset-slug>')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\n========================================================================\nROLE 3 — FEATURE EXTRACTION (v2.2) — 2.5D ADJACENT-SLICE STACKING\n========================================================================\nجديد في النسخة دي (بناءً على \"First Place Roadmap\" — بند Role 2 [P0]):\n\n  4. 2.5D Adjacent Slice Stacking بدل تكرار نفس السلايس x3:\n     كنا بناخد سلايس واحد ونكرره في الـ 3 قنوات (RGB وهمي). دلوقتي\n     بناخد 3 سلايسات حقيقية متجاورة [z-1, z, z+1] — كده الموديل بيشوف\n     سياق عمق حقيقي (امتداد تمزق أو effusion عبر السلايسات) بدل تكرار\n     نفس المعلومة. التكلفة الحسابية زيادتها صفر — نفس عدد الصور بالظبط.\n     عند حواف الـ series (أول/آخر سلايس) بيتعمل clamp للسلايس الطرفي\n     بدل ما يطلع index خارج النطاق.\n\n  5. MAX_SLICES_PER_SERIES (اختياري، OFF افتراضيًا):\n     خيار لقص كل series لأقرب N سلايس مركزية بس (بند Role 2 [P1] في\n     نفس المستند). سايبه None افتراضيًا عمدًا — ده قرار بيمس البيانات\n     نفسها، ومش هيتفعّل لحد ما تراجعه وتحدد رقم واثق فيه.\n========================================================================\nانسخ الملف ده كامل والصقه في خلية واحدة في الـ Kaggle Notebook.\n\nالفرق عن النسخة اللي فاتت (v3) — 3 تعديلات لتحسين جودة الـ embeddings\nنفسها (مش الـ MIL model)، بناءً على إن الفرق بين موديل بسيط ومعقد كان\n0.02 بس، يعني المشكلة في المعلومة جوه الـ embeddings مش في Role 4:\n\n  1. CLS + patch-mean concat بدل CLS بس:\n     كنا بناخد بس التوكن العام (x_norm_clstoken) اللي بيلخّص السلايس\n     كله في متجه واحد. دلوقتي بنضيف معاه متوسط الـ patch tokens كمان،\n     عشان تفاصيل مكانية صغيرة (تمزق محدود، كسر في مكان معين) متضيعش.\n     ده بيضاعف أبعاد الـ embedding (768 -> 1536 لو Base، 384 -> 768\n     لو Small) — التكلفة الحسابية زيادتها شبه معدومة (متوسط بسيط).\n\n  2. Per-slice z-score normalization بدل ImageNet mean/std الثابتة:\n     ImageNet stats محسوبة من صور طبيعية (كاميرات)، مش MRI. دلوقتي كل\n     سلايس بيتطبّع بالـ mean/std بتاعته هو، فالتباين الحقيقي في الإشارة\n     مبيتمسحش. ده مجاني حسابيًا (نفس عدد العمليات تقريبًا).\n\n  3. EMBEDDING_DIM بقى تلقائي:\n     الكود بيكتشف hidden_size من الموديل المتضاف فعليًا (384 لو Small،\n     768 لو Base) ويبني أبعاد الـ output تلقائيًا — يعني تقدر تبدّل بين\n     Base و Small من غير ما تلمس أي رقم في الكود.\n\nالدقة (IMAGE_SIZE) سيبتها 224 زي ما هي — رفعها لـ 336/518 بيزوّد\nالتكلفة الحسابية 2-5 أضعاف (مش مجاني زي الاتنين اللي فوق)، فسيبتها\nاختيارية: غيّر IMAGE_SIZE في CONFIG لو حبيت تجربها بعد كده لوحدها.\n\n⚠️ تنبيه: أبعاد الـ output اتغيّرت (1536 بدل 768 مع Base). أي كود\nتاني بيقرأ الـ embeddings دي بعد كده (Role 4 training script،\nrole4_mil_aggregation.yaml: embedding_dim) لازم يتحدّث بنفس القيمة.\n========================================================================\n\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nimport os\nimport sys\nfrom pathlib import Path\n\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn.functional as F\nfrom tqdm import tqdm\n\n# ------------------------------------------------------------------------\n# CONFIG\n# ------------------------------------------------------------------------\nKAGGLE_INPUT_ROOT = \"/kaggle/input\"\nOUTPUT_DIR = \"/kaggle/working\"\nDICOM_H5_HINTS = [\"dicom\", \"shard\", \"preprocess\"]\nEMBEDDINGS_HINT = \"embed\"\nMANIFEST_HINT = \"manifest\"\n\nEXTRACTION_MODE = \"cls_and_patch_mean\"   # or \"cls_only\" to go back to the old behavior\nNORM_MODE = \"per_slice_zscore\"           # or \"imagenet\" to go back to the old behavior\nIMAGE_SIZE = 224                          # bump to 336/518 separately later — not free\nSTACKING_MODE = \"adjacent_triplet\"        # or \"single_repeat\" for the old same-slice-x3 behavior\nMAX_SLICES_PER_SERIES = None              # e.g. 32 to keep only the N central slices — OFF by\n                                           # default; this changes what data gets analyzed, so\n                                           # opt in deliberately rather than as a silent default\nBATCH_SIZE = 64\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nUSE_AMP = DEVICE == \"cuda\"\nFLUSH_EVERY = 20\n\n# Only used when NORM_MODE == \"imagenet\"\nIMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406], device=DEVICE).view(1, 3, 1, 1)\nIMAGENET_STD = torch.tensor([0.229, 0.224, 0.225], device=DEVICE).view(1, 3, 1, 1)\n\n# Set after the backbone loads and its real hidden_size is known — do not\n# hardcode this, it must match whatever model (Base or Small) is attached.\nEMBEDDING_DIM = None\n\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s [%(levelname)s] %(message)s\",\n                     handlers=[logging.StreamHandler(sys.stdout)])\nlog = logging.getLogger(\"extract_features_v2_1\")\n\n\n# ========================================================================\n# PART A — SERIES FILTERING (Role 2.3) — unchanged\n# ========================================================================\nCOLUMN_ALIASES: dict[str, list[str]] = {\n    \"plane\": [\"plane\", \"Anatomical_Plane\", \"anatomical_plane\"],\n    \"fluid_sensitive\": [\"fluid_sensitive\", \"Fluid_Sensitive\"],\n    \"fat_suppression\": [\"fat_suppression\", \"Fat_Suppression\"],\n    \"sequence_type_raw\": [\"sequence_type_raw\", \"SeriesDescription\", \"sequence_type\"],\n    \"n_slices\": [\"n_slices\", \"num_slices\", \"NumSlices\", \"n_frames\"],\n}\n\n\ndef _normalize_columns(manifest: pd.DataFrame, needed: list[str]) -> pd.DataFrame:\n    manifest = manifest.copy()\n    lower_to_actual = {c.lower(): c for c in manifest.columns}\n    rename_map, still_missing = {}, []\n    for canonical in needed:\n        if canonical in manifest.columns:\n            continue\n        found = None\n        for alias in COLUMN_ALIASES.get(canonical, [canonical]):\n            if alias.lower() in lower_to_actual:\n                found = lower_to_actual[alias.lower()]\n                break\n        if found is not None:\n            rename_map[found] = canonical\n        else:\n            still_missing.append(canonical)\n    if still_missing:\n        raise ValueError(\n            f\"Manifest missing required columns: {sorted(still_missing)}. \"\n            f\"Available columns: {sorted(manifest.columns.tolist())}. \"\n            f\"Add the real name to COLUMN_ALIASES above if it's just spelled differently.\"\n        )\n    return manifest.rename(columns=rename_map)\n\n\ndef filter_series(manifest: pd.DataFrame, verbose: bool = True) -> pd.DataFrame:\n    original_count = len(manifest)\n    required = [\"StudyInstanceUID\", \"SeriesInstanceUID\", \"plane\",\n                \"fluid_sensitive\", \"fat_suppression\", \"sequence_type_raw\", \"n_slices\"]\n    manifest = _normalize_columns(manifest, required)\n\n    scout_keywords = [\"scout\", \"localizer\", \"survey\", \"loc\", \"pilot\"]\n    scout_mask = manifest[\"sequence_type_raw\"].astype(str).str.lower().str.contains(\n        \"|\".join(scout_keywords), na=False, regex=True)\n    scouts_dropped = int(scout_mask.sum())\n    manifest = manifest[~scout_mask].reset_index(drop=True)\n\n    manifest[\"fluid_sensitive\"] = pd.to_numeric(manifest[\"fluid_sensitive\"], errors=\"coerce\").fillna(0).astype(int)\n    manifest[\"n_slices\"] = pd.to_numeric(manifest[\"n_slices\"], errors=\"coerce\").fillna(0).astype(int)\n\n    missing_plane_mask = manifest[\"plane\"].isna() | (manifest[\"plane\"].astype(str).str.strip() == \"\")\n    n_missing_plane = int(missing_plane_mask.sum())\n    if n_missing_plane:\n        manifest.loc[missing_plane_mask, \"plane\"] = (\n            \"unknown_\" + manifest.loc[missing_plane_mask, \"SeriesInstanceUID\"].astype(str)\n        )\n\n    manifest = manifest.sort_values(\n        by=[\"StudyInstanceUID\", \"plane\", \"fluid_sensitive\", \"n_slices\"],\n        ascending=[True, True, False, False],\n    )\n    before_dedup = len(manifest)\n    manifest = manifest.drop_duplicates(subset=[\"StudyInstanceUID\", \"plane\"], keep=\"first\").reset_index(drop=True)\n    duplicates_dropped = before_dedup - len(manifest)\n    final_count = len(manifest)\n\n    if verbose:\n        print(f\"  [filter] original={original_count:,} scouts_dropped={scouts_dropped:,} \"\n              f\"missing_plane={n_missing_plane:,} dup_dropped={duplicates_dropped:,} \"\n              f\"final={final_count:,} ({100 * final_count / max(original_count,1):.1f}%)\")\n        if n_missing_plane:\n            print(f\"  ⚠️  {n_missing_plane} series had no plane value — check the Role 2.1 metadata join.\")\n\n    return manifest\n\n\n# ========================================================================\n# PART B — AUTO-DISCOVERY (model + shards + manifests) — unchanged\n# ========================================================================\ndef discover_model_dir(root: str = KAGGLE_INPUT_ROOT) -> tuple[str, str]:\n    candidates = []\n    for dirpath, dirnames, filenames in os.walk(root):\n        if \"dino\" in dirpath.lower():\n            candidates.append(dirpath)\n            if \"config.json\" in filenames:\n                return dirpath, \"transformers\"\n\n    for dirpath, dirnames, filenames in os.walk(root):\n        if \"dino\" not in dirpath.lower():\n            continue\n        for f in filenames:\n            if f.lower().endswith((\".pth\", \".pt\", \".bin\")):\n                return dirpath, \"raw_pytorch\"\n\n    raise FileNotFoundError(\n        \"Couldn't find an attached DINOv2 model under /kaggle/input. \"\n        f\"Directories containing 'dino' that WERE found: {candidates or 'none'}. \"\n        \"Make sure the DINO v2 model is attached via the Models panel (Add Input).\"\n    )\n\n\ndef discover_shards(root: str = KAGGLE_INPUT_ROOT) -> list[dict]:\n    all_h5, all_manifests = [], []\n    for dirpath, dirnames, filenames in os.walk(root):\n        for f in filenames:\n            full = os.path.join(dirpath, f)\n            fl = f.lower()\n            if fl.endswith(\".h5\"):\n                if EMBEDDINGS_HINT in fl:\n                    continue\n                all_h5.append(full)\n            elif fl.endswith(\".csv\") and MANIFEST_HINT in fl:\n                all_manifests.append(full)\n\n    if not all_h5:\n        raise FileNotFoundError(f\"No DICOM .h5 files found under {root}.\")\n    if not all_manifests:\n        raise FileNotFoundError(f\"No manifest CSV found under {root} (looking for '{MANIFEST_HINT}' in filename).\")\n\n    global_manifest = all_manifests[0] if len(all_manifests) == 1 else None\n\n    shards = []\n    for h5_path in sorted(all_h5):\n        h5_dir = os.path.dirname(h5_path)\n        same_dir_matches = [m for m in all_manifests if os.path.dirname(m) == h5_dir]\n        if same_dir_matches:\n            manifest_path = same_dir_matches[0]\n        elif global_manifest:\n            manifest_path = global_manifest\n        else:\n            raise FileNotFoundError(\n                f\"Can't determine which manifest belongs to {h5_path} — multiple manifests exist \"\n                f\"({all_manifests}) and none sit in the same folder.\"\n            )\n        shard_name = Path(h5_path).stem\n        shards.append({\"dicom_h5\": h5_path, \"manifest_csv\": manifest_path, \"shard_name\": shard_name})\n\n    return shards\n\n\n# ========================================================================\n# PART C — MODEL LOADING (hidden_size detected dynamically — Base or Small)\n# ========================================================================\ndef load_dinov2_model(model_dir: str, fmt: str):\n    if fmt != \"transformers\":\n        raise RuntimeError(\n            f\"Found weights at {model_dir} but no config.json — this script only \"\n            f\"knows how to load the 'transformers' framework variation of DINOv2. \"\n            f\"Attach that variation instead of the raw PyTorch one.\"\n        )\n\n    from transformers import AutoConfig, AutoModel\n\n    config = AutoConfig.from_pretrained(model_dir, local_files_only=True)\n    hidden_size = getattr(config, \"hidden_size\", None)\n    if hidden_size not in (384, 768):\n        raise ValueError(\n            f\"Unexpected hidden_size={hidden_size} — expected 384 (DINOv2 Small) or \"\n            f\"768 (DINOv2 Base). Check which variation is attached.\"\n        )\n    log.info(f\"Backbone: {model_dir} — hidden_size={hidden_size} \"\n              f\"({'Base' if hidden_size == 768 else 'Small'})\")\n\n    model = AutoModel.from_pretrained(model_dir, local_files_only=True)\n    model.eval().to(DEVICE)\n    for p in model.parameters():\n        p.requires_grad_(False)\n    return model, hidden_size\n\n\n# ========================================================================\n# PART D — BATCH PREPROCESSING + INFERENCE\n# ========================================================================\ndef _select_central_slices(slice_array_np: np.ndarray, max_slices: int | None) -> np.ndarray:\n    \"\"\"Optionally keep only the N central slices (Role 2 [P1] pruning). No-op when None.\"\"\"\n    if max_slices is None or len(slice_array_np) <= max_slices:\n        return slice_array_np\n    n = len(slice_array_np)\n    start = (n - max_slices) // 2\n    return slice_array_np[start : start + max_slices]\n\n\ndef _build_adjacent_triplets(slice_array_np: np.ndarray) -> np.ndarray:\n    \"\"\"(N, H, W) -> (N, 3, H, W) using real neighbors [z-1, z, z+1], edges clamped.\"\"\"\n    n = len(slice_array_np)\n    idx_prev = np.clip(np.arange(n) - 1, 0, n - 1)\n    idx_next = np.clip(np.arange(n) + 1, 0, n - 1)\n    return np.stack([slice_array_np[idx_prev], slice_array_np, slice_array_np[idx_next]], axis=1)\n\n\ndef prepare_and_process_batch(slice_array_np: np.ndarray, model) -> np.ndarray:\n    slice_array_np = _select_central_slices(slice_array_np, MAX_SLICES_PER_SERIES)\n    if len(slice_array_np) == 0:\n        return np.empty((0, EMBEDDING_DIM), dtype=np.float32)\n\n    if STACKING_MODE == \"adjacent_triplet\":\n        stacked = _build_adjacent_triplets(slice_array_np)  # (N, 3, H, W) — real neighbor slices\n    else:\n        stacked = np.repeat(slice_array_np[:, None, :, :], 3, axis=1)  # old behavior: same slice x3\n\n    embeddings_list = []\n    for i in range(0, len(stacked), BATCH_SIZE):\n        batch_np = stacked[i : i + BATCH_SIZE]\n        tensor = torch.from_numpy(batch_np).float().to(DEVICE)  # already (B, 3, H, W)\n\n        if tensor.shape[-2:] != (IMAGE_SIZE, IMAGE_SIZE):\n            tensor = F.interpolate(tensor, size=(IMAGE_SIZE, IMAGE_SIZE), mode=\"bilinear\", align_corners=False)\n\n        if NORM_MODE == \"per_slice_zscore\":\n            mean = tensor.mean(dim=(1, 2, 3), keepdim=True)\n            std = tensor.std(dim=(1, 2, 3), keepdim=True) + 1e-6\n            tensor = (tensor - mean) / std\n        else:  # \"imagenet\"\n            if tensor.max() > 1.0:\n                tensor = tensor / 255.0\n            tensor = (tensor - IMAGENET_MEAN) / IMAGENET_STD\n\n        with torch.inference_mode():\n            with torch.amp.autocast(\"cuda\", enabled=USE_AMP):\n                outputs = model(pixel_values=tensor)\n                cls_token = outputs.last_hidden_state[:, 0, :]\n                if EXTRACTION_MODE == \"cls_and_patch_mean\":\n                    patch_mean = outputs.last_hidden_state[:, 1:, :].mean(dim=1)\n                    features = torch.cat([cls_token, patch_mean], dim=-1)\n                else:\n                    features = cls_token\n\n        embeddings_list.append(features.float().cpu().numpy().astype(np.float32))\n\n    return np.vstack(embeddings_list)\n\n\ndef _already_done(out_h5: h5py.File, study_uid: str, series_uid: str) -> bool:\n    return study_uid in out_h5 and series_uid in out_h5[study_uid]\n\n\n# ========================================================================\n# PART E — PER-SHARD PIPELINE (resumable, fault-tolerant)\n# ========================================================================\ndef run_feature_extraction_for_shard(manifest_csv_path, dicom_h5_input_path, output_embeddings_h5, model):\n    manifest = pd.read_csv(manifest_csv_path, dtype={\"StudyInstanceUID\": str, \"SeriesInstanceUID\": str})\n    filtered_manifest = filter_series(manifest, verbose=True)\n    valid_series_set = set(filtered_manifest[\"SeriesInstanceUID\"].unique())\n\n    failures = []\n    processed_since_flush = 0\n\n    with h5py.File(dicom_h5_input_path, \"r\") as in_h5, \\\n         h5py.File(output_embeddings_h5, \"a\") as out_h5:\n\n        study_uids = list(in_h5.keys())\n        for study_uid in tqdm(study_uids, desc=f\"  {Path(output_embeddings_h5).stem}\"):\n            study_group = in_h5[study_uid]\n            out_study_group = out_h5.require_group(study_uid)\n\n            for series_uid in study_group.keys():\n                if series_uid not in valid_series_set:\n                    continue\n                if _already_done(out_h5, study_uid, series_uid):\n                    continue\n\n                try:\n                    slice_data = study_group[series_uid][:]\n                    embeddings = prepare_and_process_batch(slice_data, model)\n                    out_study_group.create_dataset(\n                        series_uid, data=embeddings, compression=\"gzip\", compression_opts=4\n                    )\n                    processed_since_flush += 1\n                except Exception as e:\n                    log.error(f\"FAILED {study_uid}/{series_uid}: {e}\")\n                    failures.append({\"StudyInstanceUID\": study_uid, \"SeriesInstanceUID\": series_uid, \"error\": str(e)})\n                    continue\n\n                if processed_since_flush >= FLUSH_EVERY:\n                    out_h5.flush()\n                    processed_since_flush = 0\n\n        out_h5.flush()\n\n    if failures:\n        failures_path = str(Path(output_embeddings_h5).with_suffix(\"\")) + \"_failures.csv\"\n        pd.DataFrame(failures).to_csv(failures_path, index=False)\n        log.warning(f\"  {len(failures)} series failed in this shard — see {failures_path}\")\n\n    with h5py.File(output_embeddings_h5, \"r\") as out_h5:\n        done_series = {s for study in out_h5.keys() for s in out_h5[study].keys()}\n    missing = valid_series_set - done_series\n    if missing:\n        log.warning(f\"  {len(missing)} series still missing — re-run this cell to resume this shard.\")\n        return False\n    return True\n\n\n# ========================================================================\n# RUN\n# ========================================================================\nmodel_dir, model_fmt = discover_model_dir()\nmodel, backbone_hidden = load_dinov2_model(model_dir, model_fmt)\n\nEMBEDDING_DIM = backbone_hidden * 2 if EXTRACTION_MODE == \"cls_and_patch_mean\" else backbone_hidden\nlog.info(f\"extraction_mode={EXTRACTION_MODE}  norm_mode={NORM_MODE}  \"\n         f\"image_size={IMAGE_SIZE}  -> output embedding_dim={EMBEDDING_DIM}\")\nlog.info(f\"⚠️  Downstream configs (role4_mil_aggregation.yaml: embedding_dim, \"\n         f\"NestedKneeMILDataset) must be updated to {EMBEDDING_DIM} to match.\")\n\nshards = discover_shards()\nlog.info(f\"Discovered {len(shards)} DICOM shard(s) to process:\")\nfor s in shards:\n    log.info(f\"  - {s['shard_name']}  (manifest: {os.path.basename(s['manifest_csv'])})\")\n\nresults = {}\nfor shard in shards:\n    out_path = os.path.join(OUTPUT_DIR, f\"{shard['shard_name']}_embeddings_v2.h5\")\n    log.info(f\"--- Processing {shard['shard_name']} -> {os.path.basename(out_path)} ---\")\n    complete = run_feature_extraction_for_shard(shard[\"manifest_csv\"], shard[\"dicom_h5\"], out_path, model)\n    results[shard[\"shard_name\"]] = \"complete\" if complete else \"incomplete — re-run to resume\"\n\nlog.info(\"========== SUMMARY ==========\")\nfor name, status in results.items():\n    log.info(f\"  {name}: {status}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}