{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":71549,"databundleVersionId":8561470,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":9245433,"datasetId":5592926,"databundleVersionId":9433882},{"sourceType":"kernelVersion","sourceId":306930178,"isSourceIdPinned":false}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ==============================================================================\n# RSNA 2024 – Lumbar Spine Degenerative Classification\n# NOTEBOOK 2 of 4: YOLO Crop Extraction + Preprocessing + DataLoaders\n# ------------------------------------------------------------------------------\n# Inputs  (add as Kaggle dataset, e.g. \"rsna-nb1-out\"):\n#   /kaggle/input/rsna-nb1-out/flat_df.pkl\n#   /kaggle/input/rsna-nb1-out/outputs/yolo_runs/lumbar/weights/best.pt\n#\n# Saves to /kaggle/working/:\n#   flat_df_with_crops.pkl\n#   preprocessed_crops.zip    (~4-8 GB of 224x224 JPEG crops)\n#\n# Upload /kaggle/working/ as a Kaggle dataset before running Notebook 3.\n!pip install ultralytics albumentations timm -q\n# ==============================================================================\n\nVIS_MODE = False\n\n# ==============================================================================\n# CELL 1: SETUP & IMPORTS (identical base to NB1)\n# ==============================================================================\nimport os, gc, math, shutil, random, warnings, time, pickle\nwarnings.filterwarnings(\"ignore\")\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport matplotlib\nif not VIS_MODE:\n    matplotlib.use('Agg')\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom scipy.ndimage import gaussian_filter\nfrom multiprocessing import Pool, cpu_count\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom ultralytics import YOLO\n\n# ── Paths ──────────────────────────────────────────────────────────────────────\nON_KAGGLE = os.path.exists('/kaggle/input')\n\nif ON_KAGGLE:\n    # ── UPDATE this to match your NB1 dataset name ────────────────────────────\n    NB1_DATASET   = '/kaggle/input/notebooks/tahuja/notebook-1'\n    BASE_DIR      = '/kaggle/input/competitions/rsna-2024-lumbar-spine-degenerative-classification'\n    PNG_DIR       = '/kaggle/input/dataset/deepakat002/rsna-lumbar-spine-test-train-png-format'\n    TRAIN_PNG_DIR = os.path.join(PNG_DIR, 'train_images_png')\n    PREP_DIR      = '/tmp/preprocessed_dataset'     # build here, zip to /kaggle/working\n    OUT_DIR       = '/kaggle/working/outputs'\n    YOLO_WEIGHTS  = f'{NB1_DATASET}/outputs/yolo_runs/lumbar/weights/best.pt'\nelse:\n    LOCAL_DATA_DIR = r'C:\\rsna'\n    NB1_DATASET    = LOCAL_DATA_DIR\n    BASE_DIR       = os.path.join(LOCAL_DATA_DIR, 'csv')\n    TRAIN_PNG_DIR  = os.path.join(LOCAL_DATA_DIR, 'png', 'train_images_png')\n    PREP_DIR       = os.path.join(LOCAL_DATA_DIR, 'preprocessed_dataset')\n    OUT_DIR        = os.path.join(LOCAL_DATA_DIR, 'outputs')\n    YOLO_WEIGHTS   = os.path.join(LOCAL_DATA_DIR, 'outputs', 'yolo_runs', 'lumbar', 'weights', 'best.pt')\n\nfor d in [PREP_DIR, OUT_DIR]:\n    os.makedirs(d, exist_ok=True)\n\n# ── Hyper-parameters ───────────────────────────────────────────────────────────\nIMG_SIZE     = 224\nSEED         = 42\nN_FOLDS      = 5\nFOLD         = 0\nYOLO_CONF    = 0.25\nYOLO_PAD     = 0.15\nDEVICE       = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# NB2 has no GPU training — only YOLO inference (CPU-bound crop loop is parallelized)\nNUM_WORKERS  = min(4, cpu_count()) if ON_KAGGLE else 2\n# DataLoader workers for NB2 val check at end\nDL_WORKERS   = 4 if ON_KAGGLE else 2\n\nSEVERITY_MAP = {'Normal/Mild': 0, 'Moderate': 1, 'Severe': 2}\nSEV_INV      = {0: 'Normal/Mild', 1: 'Moderate', 2: 'Severe'}\nLEVELS       = ['l1_l2', 'l2_l3', 'l3_l4', 'l4_l5', 'l5_s1']\nLEVEL_TO_CLS = {lv: i for i, lv in enumerate(LEVELS)}\nCLS_TO_LEVEL = {i: lv for lv, i in LEVEL_TO_CLS.items()}\nSEV_COLORS   = {0: '#4CAF50', 1: '#FF9800', 2: '#F44336'}\n\ndef seed_all(s=SEED):\n    random.seed(s); np.random.seed(s)\n    torch.manual_seed(s); torch.cuda.manual_seed_all(s)\n\nseed_all()\nprint(f\"Device      : {DEVICE}\")\nprint(f\"NUM_WORKERS : {NUM_WORKERS}\")\nprint(f\"✅ Cell 1 Complete\")\n\n\n# ==============================================================================\n# CELL 2: LOAD flat_df FROM NB1\n# ==============================================================================\nprint(\"\\n\" + \"=\"*65)\nprint(\"  CELL 2 — Load flat_df from NB1\")\nprint(\"=\"*65)\n\nflat_df = pd.read_pickle(f'{NB1_DATASET}/flat_df.pkl')\nprint(f\"Loaded flat_df: {flat_df.shape}\")\nprint(flat_df['severity_str'].value_counts().to_string())\nprint(\"✅ Cell 2 Complete\")\n\n\n# ==============================================================================\n# CELL 3: PREPROCESSING FUNCTIONS (same as NB1 Cell 3)\n# ==============================================================================\nprint(\"\\n\" + \"=\"*65)\nprint(\"  CELL 3 — Preprocessing functions\")\nprint(\"=\"*65)\n\ndef step_minmax_norm(img):\n    mn, mx = img.min(), img.max()\n    if mx - mn < 1e-6: return np.zeros_like(img, dtype=np.uint8)\n    return (255.0 * (img - mn) / (mx - mn)).astype(np.uint8)\n\ndef step_windowing(img, low_pct=1.0, high_pct=99.0):\n    lo = np.percentile(img, low_pct); hi = np.percentile(img, high_pct)\n    return (255.0 * (np.clip(img, lo, hi) - lo) / max(hi - lo, 1e-6)).astype(np.uint8)\n\ndef step_denoise(img, sigma=1.0):\n    \"\"\"Gaussian — replaces NL-means for speed.\"\"\"\n    return cv2.GaussianBlur(img, (0, 0), sigma)\n\ndef step_clahe(img, clip_limit=2.0, tile_grid=(8, 8)):\n    return cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid).apply(img)\n\ndef step_unsharp_mask(img, sigma=1.0, strength=1.5):\n    blurred = gaussian_filter(img.astype(np.float32), sigma=sigma)\n    return np.clip(img.astype(np.float32) + strength * (img.astype(np.float32) - blurred), 0, 255).astype(np.uint8)\n\ndef step_gamma_correction(img, gamma=None):\n    if gamma is None:\n        mean_val = img.mean()\n        gamma = float(np.clip(math.log(128.0/255.0) / math.log(max(mean_val,1.0)/255.0 + 1e-6), 0.4, 2.5))\n    lut = np.array([((i/255.0)**(1.0/gamma))*255 for i in range(256)], dtype=np.uint8)\n    return cv2.LUT(img, lut)\n\ndef step_background_mask(img, morph_iter=3):\n    _, mask = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n    kernel  = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))\n    mask    = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=morph_iter)\n    return cv2.bitwise_and(img, img, mask=mask)\n\ndef step_to_rgb(img_gray):\n    g = img_gray.copy()\n    if g.dtype != np.uint8:\n        mn, mx = g.min(), g.max()\n        g = (255.0 * (g - mn) / max(mx - mn, 1e-6)).astype(np.uint8)\n    return cv2.cvtColor(g, cv2.COLOR_GRAY2RGB)\n\ndef preprocess_mri_array(img_bgr: np.ndarray) -> np.ndarray:\n    \"\"\"Preprocess a BGR crop array → uint8 RGB.\"\"\"\n    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)\n    img = step_minmax_norm(img)\n    img = step_windowing(img)\n    img = step_denoise(img)\n    img = step_clahe(img)\n    img = step_unsharp_mask(img)\n    img = step_gamma_correction(img)\n    img = step_background_mask(img)\n    return step_to_rgb(img)\n\nprint(\"✅ Cell 3 Complete — preprocessing functions defined\")\n\n\n# ==============================================================================\n# CELL 4: YOLO CROP EXTRACTION — PARALLELIZED\n# ==============================================================================\nprint(\"\\n\" + \"=\"*65)\nprint(\"  CELL 4 — YOLO Crop Extraction (parallel, JPEG output)\")\nprint(\"=\"*65)\n\ndef do_crop(img_bgr: np.ndarray, x1, y1, x2, y2, pad=YOLO_PAD) -> np.ndarray:\n    h, w   = img_bgr.shape[:2]\n    px     = (x2 - x1) * pad; py = (y2 - y1) * pad\n    x1 = max(0, int(x1-px)); y1 = max(0, int(y1-py))\n    x2 = min(w, int(x2+px)); y2 = min(h, int(y2+py))\n    return img_bgr[y1:y2, x1:x2]\n\n\n# ── Worker function (each worker loads its own YOLO instance) ─────────────────\ndef process_one_row(args):\n    \"\"\"\n    Top-level function for multiprocessing.Pool.\n    Each worker gets its own YOLO detector to avoid sharing CUDA contexts.\n    Returns list of crop record dicts.\n    \"\"\"\n    (row_dict, PREP_DIR, YOLO_WEIGHTS, YOLO_CONF, YOLO_PAD,\n     CLS_TO_LEVEL, LEVELS) = args\n\n    records = []\n    try:\n        detector = YOLO(YOLO_WEIGHTS)   # loaded fresh per worker process\n        img_path = row_dict['img_path']\n        img_bgr  = cv2.imread(img_path)\n        if img_bgr is None:\n            return records\n\n        results = detector(img_path, conf=YOLO_CONF, iou=0.45,\n                           verbose=False, device='cpu')[0]   # CPU inference in workers\n\n        crops_to_save = []\n        if len(results.boxes) == 0:\n            h, w      = img_bgr.shape[:2]\n            fallback  = img_bgr[h//4:h*3//4, w//4:w*3//4]\n            for lv in LEVELS:\n                crops_to_save.append((fallback, lv, 0.0))\n        else:\n            for box in results.boxes:\n                cls_id = int(box.cls.item())\n                level  = CLS_TO_LEVEL.get(cls_id)\n                if level is None: continue\n                x1, y1, x2, y2 = box.xyxy[0].tolist()\n                crop = do_crop(img_bgr, x1, y1, x2, y2, pad=YOLO_PAD)\n                if crop.size == 0: continue\n                crops_to_save.append((crop, level, float(box.conf.item())))\n\n        for crop_bgr, level, conf in crops_to_save:\n            try:\n                crop_prep = preprocess_mri_array(crop_bgr)\n            except Exception:\n                crop_prep = cv2.cvtColor(\n                    cv2.resize(crop_bgr, (224, 224)), cv2.COLOR_BGR2RGB)\n\n            # ── JPEG output — ~60% smaller than PNG, no meaningful quality loss\n            fname = (f\"{int(row_dict['study_id'])}_{int(row_dict['series_id'])}_\"\n                     f\"{int(row_dict['instance_number'])}_{level}.jpg\")\n            dst   = os.path.join(PREP_DIR, fname)\n            cv2.imwrite(dst,\n                        cv2.cvtColor(crop_prep, cv2.COLOR_RGB2BGR),\n                        [cv2.IMWRITE_JPEG_QUALITY, 95])\n\n            records.append({\n                'study_id':        row_dict['study_id'],\n                'series_id':       row_dict['series_id'],\n                'instance_number': row_dict['instance_number'],\n                'img_path':        img_path,\n                'level_norm':      level,\n                'crop_path':       dst,\n                'det_conf':        conf,\n            })\n    except Exception as e:\n        pass   # silently skip corrupt images\n    return records\n\n\n# ── Build arg list ─────────────────────────────────────────────────────────────\nunique_for_crop = (flat_df[['study_id', 'series_id', 'instance_number', 'img_path']]\n                   .drop_duplicates()\n                   .reset_index(drop=True))\n\narg_list = [\n    (row.to_dict(), PREP_DIR, YOLO_WEIGHTS, YOLO_CONF, YOLO_PAD,\n     CLS_TO_LEVEL, LEVELS)\n    for _, row in unique_for_crop.iterrows()\n]\n\nprint(f\"Processing {len(arg_list)} unique images with {NUM_WORKERS} workers...\")\n\n# ── Parallelized crop loop ─────────────────────────────────────────────────────\n# Note: YOLO runs on CPU inside workers — GPU is reserved for NB3 training.\n# Pool.imap_unordered gives progress + early error visibility.\nall_records = []\nwith Pool(processes=NUM_WORKERS) as pool:\n    for batch in tqdm(pool.imap_unordered(process_one_row, arg_list),\n                      total=len(arg_list), desc='Crop + Preprocess'):\n        all_records.extend(batch)\n\ncrop_df = pd.DataFrame(all_records)\nprint(f\"Total crops saved: {len(crop_df)}\")\n\n# ── Merge crop paths back into flat_df ────────────────────────────────────────\nflat_df = flat_df.merge(\n    crop_df[['study_id', 'series_id', 'instance_number', 'level_norm', 'crop_path', 'det_conf']],\n    on=['study_id', 'series_id', 'instance_number', 'level_norm'],\n    how='left')\nflat_df['crop_path'] = flat_df['crop_path'].fillna(flat_df['img_path'])\n\nif VIS_MODE:\n    n_show = min(5, len(crop_df))\n    sample_rows = (crop_df[crop_df['det_conf'] > 0]\n                   .groupby('level_norm').first().reset_index().head(n_show))\n    fig, axes = plt.subplots(2, n_show, figsize=(n_show*3.5, 7))\n    fig.patch.set_facecolor('#0D1117')\n    for col, (_, cr) in enumerate(sample_rows.iterrows()):\n        orig_bgr = cv2.imread(cr['img_path'])\n        prep_bgr = cv2.imread(cr['crop_path'])\n        for row_i, (bgr, title, color) in enumerate([\n            (orig_bgr, f\"Original\\n{cr['level_norm']}\", '#8B949E'),\n            (prep_bgr, f\"Preprocessed\\nconf={cr['det_conf']:.2f}\", '#4CAF50'),\n        ]):\n            ax = axes[row_i][col]; ax.set_facecolor('#161B22'); ax.axis('off')\n            if bgr is not None:\n                ax.imshow(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB), aspect='auto')\n            ax.set_title(title, color=color, fontsize=8)\n    plt.suptitle('CELL 4 — YOLO Crops: Raw vs Preprocessed', color='white', fontsize=13, fontweight='bold')\n    plt.tight_layout()\n    plt.savefig(f'{OUT_DIR}/cell4_crops.png', dpi=130, bbox_inches='tight', facecolor='#0D1117')\n    plt.show()\n\nprint(f\"✅ Cell 4 Complete — {len(crop_df)} crops in {PREP_DIR}\")\n\n\n# ==============================================================================\n# CELL 5: AUGMENTATION PIPELINE + DATALOADERS\n# ==============================================================================\nprint(\"\\n\" + \"=\"*65)\nprint(\"  CELL 5 — Augmentation + DataLoaders\")\nprint(\"=\"*65)\n\nTRAIN_TF = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.HorizontalFlip(p=0.5),\n    A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=10, p=0.4),\n    A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.3),\n    A.GaussNoise(var_limit=(5, 25), p=0.2),\n    A.CoarseDropout(max_holes=4, max_height=16, max_width=16, min_holes=1, p=0.2),\n    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ToTensorV2(),\n])\nVAL_TF = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ToTensorV2(),\n])\n\nclass RSNADataset(Dataset):\n    def __init__(self, df: pd.DataFrame, transform=None):\n        self.df = df.dropna(subset=['crop_path']).reset_index(drop=True)\n        self.tf = transform\n\n    def __len__(self): return len(self.df)\n\n    def __getitem__(self, idx):\n        row   = self.df.iloc[idx]\n        label = int(row['severity'])\n        img   = cv2.imread(str(row['crop_path']))\n        if img is None:\n            img = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        if self.tf:\n            img = self.tf(image=img)['image']\n        return img.float(), torch.tensor(label, dtype=torch.long)\n\n# ── StratifiedGroupKFold ──────────────────────────────────────────────────────\nsgkf   = StratifiedGroupKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\nstrat  = flat_df['severity'].astype(str) + '_' + flat_df['condition_norm']\ngroups = flat_df['study_id'].values\n\nfor fold_idx, (tri, vli) in enumerate(sgkf.split(flat_df, strat, groups)):\n    if fold_idx == FOLD:\n        trn_df = flat_df.iloc[tri].reset_index(drop=True)\n        val_df = flat_df.iloc[vli].reset_index(drop=True)\n        break\n\n# ── Class weights ─────────────────────────────────────────────────────────────\ntrn_labels  = trn_df['severity'].values.astype(int)\ncls_counts  = np.bincount(trn_labels, minlength=3).astype(float)\ncls_counts  = np.where(cls_counts == 0, 1, cls_counts)\ncls_weights = len(trn_labels) / (3 * cls_counts)\nprint(f\"Fold {FOLD}  train={len(trn_df)}  val={len(val_df)}\")\nprint(f\"Class weights: {dict(zip(SEV_INV.values(), cls_weights.round(3)))}\")\n\n# Quick sanity check — build loaders and verify one batch loads\nfrom torch.utils.data import WeightedRandomSampler\ntrain_dataset = RSNADataset(trn_df, TRAIN_TF)\nval_dataset   = RSNADataset(val_df, VAL_TF)\nsample_w      = torch.from_numpy(cls_weights[trn_labels]).float()\nsampler       = WeightedRandomSampler(sample_w, len(trn_labels), replacement=True)\n\ntrain_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler,\n                          num_workers=DL_WORKERS, pin_memory=True,\n                          drop_last=True, persistent_workers=True,\n                          prefetch_factor=2, timeout=60)\nval_loader   = DataLoader(val_dataset,   batch_size=64, shuffle=False,\n                          num_workers=DL_WORKERS, pin_memory=True,\n                          persistent_workers=True, prefetch_factor=2, timeout=60)\n\nimgs_b, labs_b = next(iter(train_loader))\nprint(f\"Batch check — imgs: {imgs_b.shape}  labels: {labs_b.shape}\")\n\nprint(\"✅ Cell 5 Complete\")\n\n\n# ==============================================================================\n# CELL 6: SAVE OUTPUTS FOR NOTEBOOK 3\n# ==============================================================================\nprint(\"\\n\" + \"=\"*65)\nprint(\"  CELL 6 — Save outputs for NB3\")\nprint(\"=\"*65)\n\n# Save flat_df with crop paths\nflat_df.to_pickle('/kaggle/working/flat_df_with_crops.pkl')\nprint(f\"Saved flat_df_with_crops.pkl  ({len(flat_df)} rows)\")\n\n# Save class weights so NB3 can use them directly\nnp.save('/kaggle/working/cls_weights.npy', cls_weights)\nprint(f\"Saved cls_weights.npy\")\n\n# ── Zip the preprocessed crops folder → single file, avoids 500-file dataset cap\nprint(\"Zipping preprocessed crops (this takes a few minutes)...\")\nzip_path = '/kaggle/working/preprocessed_crops'\nshutil.make_archive(zip_path, 'zip', PREP_DIR)\nzip_size = os.path.getsize(zip_path + '.zip') / 1e9\nprint(f\"Zipped crops → {zip_path}.zip  ({zip_size:.2f} GB)\")\n\nprint(\"\\n  ✅ Notebook 2 Complete!\")\nprint(\"  Upload /kaggle/working/ contents as Kaggle dataset before running NB3.\")\nprint(\"  Required files:\")\nprint(\"    flat_df_with_crops.pkl\")\nprint(\"    cls_weights.npy\")\nprint(\"    preprocessed_crops.zip\")","metadata":{"_uuid":"577ada67-a813-45c4-871f-50f2a8f12f07","_cell_guid":"70e548cb-da0c-4e18-9658-23ad57eba0d4","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}