{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "xvr-colab.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Register your own X-ray to your own CT\n",
    "\n",
    "This notebook runs [`xvr`](https://github.com/eigenvivek/xvr) on data you provide. `xvr` is the reference implementation of Gopalakrishnan and colleagues, *Rapid patient-specific neural networks for X-ray to volume registration*, Nature 2026 ([doi:10.1038/s41586-026-11045-x](https://doi.org/10.1038/s41586-026-11045-x)), MIT licensed. The differentiable renderer underneath it is [`nanodrr`](https://github.com/eigenvivek/nanodrr), by the same author.\n",
    "\n",
    "The [browser demo](https://vanioantunes.com/3d/xvr/) this notebook accompanies shows the idea on a fixed CT. This runs the real thing on yours.\n",
    "\n",
    "## Before you start\n",
    "\n",
    "**Set a GPU runtime.** *Runtime \u2192 Change runtime type \u2192 T4 GPU*. It works on CPU, slowly.\n",
    "\n",
    "**What you need ready.**\n",
    "\n",
    "1. A CT or MR volume as NIfTI, `.nii` or `.nii.gz`. If you have a DICOM folder, cell 3 converts it.\n",
    "2. One radiograph or fluoroscopy frame, as DICOM or as a 16-bit image.\n",
    "3. The geometry of the machine that took the radiograph: source-to-detector distance in mm, detector pixel spacing in mm, and the image size in pixels. In a DICOM these are `DistanceSourceToDetector` (0018,1110), `ImagerPixelSpacing` (0018,1164) and `Rows`/`Columns`.\n",
    "\n",
    "**Privacy.** Colab runs on Google's machines. Anything uploaded here leaves your computer. Do not put identifiable patient data in it without the approvals that requires. Research and teaching only; this is not a medical device and nothing here is validated for clinical use.\n",
    "\n",
    "**Registration needs a starting guess.** `xvr register` is a local optimiser, exactly like the browser demo: it refines a pose, it does not find one from nowhere. Either give it an initial pose in roughly the right neighbourhood, or train a patient-specific network first with `xvr train` so the network supplies one. Both are below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Check the runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi || echo \"No GPU. Runtime > Change runtime type > T4 GPU, then run this cell again.\"\n",
    "\n",
    "import sys\n",
    "import torch\n",
    "\n",
    "print(\"python\", sys.version.split()[0])\n",
    "print(\"torch\", torch.__version__, \"cuda\", torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Install xvr\n",
    "\n",
    "One package, from PyPI. A minute or two on a fresh runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q xvreg\n",
    "\n",
    "!xvr --version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Upload your data\n",
    "\n",
    "Run the cell and pick your files. Upload the CT (`.nii` or `.nii.gz`) and the radiograph together.\n",
    "\n",
    "For anything larger than a few hundred megabytes, mount Drive instead: uncomment the Drive lines and put the files in a folder there. Browser uploads of large NIfTI files are slow and time out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "WORK = Path(\"/content/xvr-run\")\n",
    "WORK.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Option A: upload from your computer.\n",
    "from google.colab import files\n",
    "\n",
    "uploaded = files.upload()\n",
    "for name in uploaded:\n",
    "    Path(name).rename(WORK / name)\n",
    "\n",
    "# Option B: use Google Drive instead of uploading.\n",
    "# from google.colab import drive\n",
    "# drive.mount(\"/content/drive\")\n",
    "# WORK = Path(\"/content/drive/MyDrive/xvr-run\")\n",
    "\n",
    "for f in sorted(WORK.iterdir()):\n",
    "    print(f\"{f.name:45s} {f.stat().st_size / 1e6:8.1f} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### If your CT is a DICOM folder\n",
    "\n",
    "Zip the folder, upload the zip above, then run this cell to unpack and convert it to a single NIfTI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q dicom2nifti\n",
    "\n",
    "import zipfile\n",
    "\n",
    "import dicom2nifti\n",
    "\n",
    "DICOM_ZIP = WORK / \"ct-dicom.zip\"   # change to your file name\n",
    "\n",
    "if DICOM_ZIP.exists():\n",
    "    dicom_dir = WORK / \"ct-dicom\"\n",
    "    with zipfile.ZipFile(DICOM_ZIP) as z:\n",
    "        z.extractall(dicom_dir)\n",
    "    out = WORK / \"nifti\"\n",
    "    out.mkdir(exist_ok=True)\n",
    "    dicom2nifti.convert_directory(str(dicom_dir), str(out), compression=True)\n",
    "    for f in sorted(out.glob(\"*.nii.gz\")):\n",
    "        print(\"wrote\", f)\n",
    "else:\n",
    "    print(f\"No {DICOM_ZIP.name}; skip this cell if your CT is already NIfTI.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Read the geometry out of the radiograph\n",
    "\n",
    "If the radiograph is a DICOM, this pulls the three numbers `xvr` needs. If it is not, read them off the console of the machine that took it and fill them in by hand in cell 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q pydicom\n",
    "\n",
    "import pydicom\n",
    "\n",
    "XRAY = WORK / \"xray.dcm\"   # change to your file name\n",
    "\n",
    "if XRAY.exists() and XRAY.suffix.lower() == \".dcm\":\n",
    "    ds = pydicom.dcmread(str(XRAY))\n",
    "    print(\"source to detector (mm):\", getattr(ds, \"DistanceSourceToDetector\", \"not in header\"))\n",
    "    print(\"source to patient (mm): \", getattr(ds, \"DistanceSourceToPatient\", \"not in header\"))\n",
    "    print(\"pixel spacing (mm):     \", getattr(ds, \"ImagerPixelSpacing\", getattr(ds, \"PixelSpacing\", \"not in header\")))\n",
    "    print(\"size (rows x cols):     \", ds.Rows, \"x\", ds.Columns)\n",
    "else:\n",
    "    print(\"Not a DICOM, or the name above is wrong. Fill the numbers in by hand in cell 6.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. What the CLI expects\n",
    "\n",
    "Flag names change between `xvr` releases, so read them from the version you just installed rather than trusting a notebook written months ago. Run this, then build the command in the next cell from what it prints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!xvr register --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Register\n",
    "\n",
    "Set your file names and geometry, then run. If `--help` above names the arguments differently, edit the command to match.\n",
    "\n",
    "Preprocessing matters more than people expect. A DRR is bright where the patient is dense, and some detectors store the opposite; if the optimiser goes nowhere, check that your radiograph and a rendered DRR have the same polarity before blaming the optimiser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shlex\n",
    "import subprocess\n",
    "\n",
    "VOLUME = WORK / \"ct.nii.gz\"        # your CT\n",
    "XRAY = WORK / \"xray.dcm\"           # your radiograph\n",
    "OUTDIR = WORK / \"registration\"\n",
    "OUTDIR.mkdir(exist_ok=True)\n",
    "\n",
    "SDD = 1020.0     # source to detector, mm\n",
    "DELX = 0.194     # detector pixel spacing, mm\n",
    "HEIGHT = 1536    # image height, pixels\n",
    "\n",
    "# A starting pose in roughly the right neighbourhood: rotations in degrees,\n",
    "# translations in mm, in the convention xvr prints under --help.\n",
    "INIT_ROT = (0.0, 0.0, 0.0)\n",
    "INIT_XYZ = (0.0, 0.0, 1020.0)\n",
    "\n",
    "command = f\"\"\"\n",
    "xvr register {shlex.quote(str(XRAY))}\n",
    "  -v {shlex.quote(str(VOLUME))}\n",
    "  -o {shlex.quote(str(OUTDIR))}\n",
    "  --sdd {SDD} --delx {DELX} --height {HEIGHT}\n",
    "  --rot {INIT_ROT[0]} {INIT_ROT[1]} {INIT_ROT[2]}\n",
    "  --xyz {INIT_XYZ[0]} {INIT_XYZ[1]} {INIT_XYZ[2]}\n",
    "\"\"\"\n",
    "command = \" \".join(command.split())\n",
    "print(command, \"\\n\")\n",
    "\n",
    "result = subprocess.run(command, shell=True, capture_output=True, text=True)\n",
    "print(result.stdout)\n",
    "if result.returncode != 0:\n",
    "    print(\"--- failed, stderr follows ---\")\n",
    "    print(result.stderr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Train a patient-specific network first\n",
    "\n",
    "This is the part the browser demo cannot do, and the reason the paper exists. `xvr train` learns a pose regressor from your CT alone, by rendering it from thousands of sampled poses. A handful of minutes on a T4. Afterwards registration needs no initial guess from you: the network supplies one and the optimiser refines it.\n",
    "\n",
    "Read the flags first, for the same reason as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!xvr train --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Look at the result\n",
    "\n",
    "Registration is only believable if you look at it. Overlay the radiograph and a DRR rendered at the recovered pose: where the two agree the overlay goes grey, and where they do not you get colour fringes, exactly as in the browser demo.\n",
    "\n",
    "This lists what `xvr` wrote, and gives you the overlay function to plot with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "for f in sorted(OUTDIR.rglob(\"*\")):\n",
    "    if f.is_file():\n",
    "        print(f.relative_to(OUTDIR), f\"{f.stat().st_size / 1e3:.0f} kB\")\n",
    "\n",
    "\n",
    "def overlay(fixed, moving, title=\"fixed in red, moving in cyan, grey where they agree\"):\n",
    "    \"\"\"Both arrays are 2D. Each is windowed to its own 1st and 99.5th percentile.\"\"\"\n",
    "\n",
    "    def window(a):\n",
    "        a = np.asarray(a, dtype=np.float32)\n",
    "        lo, hi = np.percentile(a, [1.0, 99.5])\n",
    "        return np.clip((a - lo) / max(hi - lo, 1e-6), 0, 1)\n",
    "\n",
    "    rgb = np.stack([window(fixed), window(moving), window(moving)], axis=-1)\n",
    "    plt.figure(figsize=(6, 6))\n",
    "    plt.imshow(rgb)\n",
    "    plt.title(title, fontsize=9)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# overlay(fixed_array, drr_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Where to go next\n",
    "\n",
    "- Full CLI documentation: [xvr.csail.mit.edu](https://xvr.csail.mit.edu/)\n",
    "- Pretrained models and the datasets from the paper, on HuggingFace: `eigenvivek/xvr` and `eigenvivek/xvr-data`\n",
    "- The renderer on its own, to build something else on the forward model: [`nanodrr`](https://github.com/eigenvivek/nanodrr)\n",
    "- The browser version of the optimisation step, with no install: [vanioantunes.com/3d/xvr/](https://vanioantunes.com/3d/xvr/)\n",
    "\n",
    "Notebook by Vanio Antunes. `xvr` and `nanodrr` are by Vivek Gopalakrishnan and colleagues, MIT licensed; this notebook is a wrapper around them and claims no part of that work."
   ]
  }
 ]
}