1-feature-extraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5cbd0b47",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c66263",
   "metadata": {},
   "source": [
    "## built-in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "2dd382f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c03327",
   "metadata": {},
   "source": [
    "## standard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "efa0518f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fce460c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "device = 'cuda:1'\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c95d6c5",
   "metadata": {},
   "source": [
    "## local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "361a548d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from source.constants import ALL_CANCER_TYPES\n",
    "from source.constants import ALL_IMG_NORMS, ALL_EXTRACTOR_MODELS\n",
    "from source.constants import DATASET_SPECIFIC_NORMALIZATION_CONSTANTS_PATH\n",
    "from source.constants import DATA_DIR, FEATURE_VECTORS_SAVE_DIR\n",
    "\n",
    "print(f\"DATA_DIR: {DATA_DIR}\")\n",
    "print(f\"FEATURE_VECTORS_SAVE_DIR: {FEATURE_VECTORS_SAVE_DIR}\")\n",
    "\n",
    "print(f\"ALL_CANCER_TYPES: {ALL_CANCER_TYPES}\")\n",
    "print(f\"ALL_IMG_NORMS: {ALL_IMG_NORMS}\")\n",
    "print(f\"DATASET_SPECIFIC_NORMALIZATION_CONSTANTS_PATH: {DATASET_SPECIFIC_NORMALIZATION_CONSTANTS_PATH}\")\n",
    "print(f\"ALL_EXTRACTOR_MODELS: {ALL_EXTRACTOR_MODELS}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a9a9e4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from source.feature_extraction.data import get_data_transform\n",
    "# help(get_data_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f9f04428",
   "metadata": {},
   "outputs": [],
   "source": [
    "from extract_features import (\n",
    "    prepare_directories,\n",
    "    calculate_dataset_mean_std,\n",
    "    update_dataset_specific_mean_std,\n",
    "    make_pytorch_dataset,\n",
    "    make_pytorch_dataloader,\n",
    "    prepare_feature_extractor,\n",
    "    extract_features,\n",
    "    save_features,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eb659c0",
   "metadata": {},
   "source": [
    "## autoreload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c21d525",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eda37eac",
   "metadata": {},
   "source": [
    "# Notebook-level Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b213c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "CANCER_TYPE = 'lung_aca'\n",
    "\n",
    "EXTRACTOR_NAME = \"owkin-phikon\"\n",
    "IMG_NORM = \"imagenet\"\n",
    "if IMG_NORM.startswith(\"lc25k\"):\n",
    "    assert CANCER_TYPE in IMG_NORM\n",
    "\n",
    "\n",
    "BATCH_SIZE = 256\n",
    "\n",
    "assert CANCER_TYPE in ALL_CANCER_TYPES\n",
    "assert IMG_NORM in ALL_IMG_NORMS\n",
    "assert EXTRACTOR_NAME in ALL_EXTRACTOR_MODELS\n",
    "\n",
    "os.listdir(DATA_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f33845a",
   "metadata": {},
   "source": [
    "# Prepare location to save features, ids, and ids_2_img_paths mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f3f807c",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dir, features_save_dir = prepare_directories(\n",
    "    all_img_dir_path=DATA_DIR,\n",
    "    all_features_save_dir=FEATURE_VECTORS_SAVE_DIR,\n",
    "    cancer_type=CANCER_TYPE,\n",
    "    img_norm=IMG_NORM,\n",
    "    extractor_name=EXTRACTOR_NAME,\n",
    ")\n",
    "print(f\"img_dir:\\n {img_dir}\")\n",
    "print(f\"features_save_dir:\\n {features_save_dir}\")\n",
    "\n",
    "features_save_paths = {\n",
    "    'ids': f'{features_save_dir}/ids.npy',\n",
    "    'ids_2_img_paths': f'{features_save_dir}/ids_2_img_paths.json',\n",
    "    'features': f'{features_save_dir}/features.npy'\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959a8888",
   "metadata": {},
   "source": [
    "# Get Data Transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "458565b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    data_transform = get_data_transform(img_norm=IMG_NORM)\n",
    "except KeyError as e:\n",
    "    print(f\"Key {e} not found in either constansts_zoo of `data.get_norm_constants()` or data-specific transforms in {DATASET_SPECIFIC_NORMALIZATION_CONSTANTS_PATH}\")\n",
    "    print(\"Calculating mean and std for the dataset...\")\n",
    "    mean, std = calculate_dataset_mean_std(img_dir=img_dir, batch_size=BATCH_SIZE)\n",
    "    data_transform = get_data_transform(img_norm='manual', mean=mean, std=std)\n",
    "    update_dataset_specific_mean_std(json_path=DATASET_SPECIFIC_NORMALIZATION_CONSTANTS_PATH, img_norm=IMG_NORM, mean=mean, std=std)\n",
    "\n",
    "print(CANCER_TYPE)\n",
    "print(data_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05d91f80",
   "metadata": {},
   "source": [
    "# Get Feature Extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48f59dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(EXTRACTOR_NAME)\n",
    "feature_extractor = prepare_feature_extractor(extractor_name=EXTRACTOR_NAME, device=device)\n",
    "feature_extractor.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c8e38b",
   "metadata": {},
   "source": [
    "# Initialise a Dataset Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a71458f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = make_pytorch_dataset(img_dir=img_dir, data_transform=data_transform)\n",
    "dataset.__getitem__(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a280c0c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.img_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc3c526d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ouput = feature_extractor(dataset.__getitem__(0)[\"image\"].to(device).unsqueeze(0))\n",
    "ouput"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5bcc2ac",
   "metadata": {},
   "source": [
    "# Initialise a Dataloader Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1eec9bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = make_pytorch_dataloader(\n",
    "    dataset=dataset, batch_size=BATCH_SIZE)\n",
    "print(\"Total instances: \", len(dataloader.dataset))\n",
    "print(\"Total batches: \", len(dataloader))\n",
    "print()\n",
    "\n",
    "first_batch = next(iter(dataloader))\n",
    "for key, val in first_batch.items():\n",
    "    if isinstance(val, torch.Tensor):\n",
    "        print(key, type(val), \":\", val.shape)\n",
    "    elif isinstance(val, list):\n",
    "        print(key, type(val), \":\", len(val))\n",
    "    else:\n",
    "        print(key, \":\", type(val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa401580",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):\n",
    "    ouput = feature_extractor(first_batch[\"image\"].to(device))\n",
    "\n",
    "print(ouput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f41a5968",
   "metadata": {},
   "source": [
    "# Run inference on the whole dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50b2b963",
   "metadata": {},
   "outputs": [],
   "source": [
    "features_and_info = extract_features(\n",
    "    feature_extractor=feature_extractor,\n",
    "    dataloader=dataloader,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5c3a746",
   "metadata": {},
   "source": [
    "# Save features, ids, and ids_2_img_paths mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "395a0b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_features(contents=features_and_info, paths=features_save_paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52d3519f",
   "metadata": {},
   "source": [
    "# Load saved features, ids, and ids_2_img_paths mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a7046d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(features_save_paths[\"ids\"])\n",
    "print(np.load(features_save_paths[\"ids\"]))\n",
    "\n",
    "print(features_save_paths[\"features\"])\n",
    "print(np.load(features_save_paths[\"features\"])[0])\n",
    "\n",
    "print(features_save_paths[\"ids_2_img_paths\"])\n",
    "with open(features_save_paths[\"ids_2_img_paths\"], \"r\") as f:\n",
    "    ids_2_img_paths = json.load(f)\n",
    "print(ids_2_img_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddafa094",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lc25k-cleaning",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}