Upload 4 files

requirements.txt

@@ -0,0 +1,6 @@
+torch>=1.9.0
+torchvision>=0.10.0
+scikit-learn>=0.24.2
+pillow>=8.0.0
+numpy>=1.19.5
+pyradiomics>=3.0.1

run_sample.ipynb

@@ -0,0 +1,883 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/javascript": "\nvar cell = this.closest('.cell');\nif (cell) {\n    cell.classList.remove('output_scroll');\n}\n",
+      "text/plain": [
+       "<IPython.core.display.Javascript object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Bad key paths in file /export/home/daifang/.config/matplotlib/matplotlibrc, line 3 ('paths: /export/home/daifang/fonts/arial/')\n",
+      "You probably need to get an updated matplotlibrc file from\n",
+      "https://github.com/matplotlib/matplotlib/blob/v3.3.4/matplotlibrc.template\n",
+      "or from the matplotlib source distribution\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda:1\n",
+      "\n",
+      "[Epoch 001] Train Loss=4.903 | Val Loss=4.777\n",
+      "  All     | Subtype AUC=0.610 | TNM AUC=0.571 | DFS C-index=0.631 | OS C-index=0.454\n",
+      "  Immune  | Subtype AUC=0.571 | TNM AUC=0.585 | DFS C-index=0.627 | OS C-index=0.478\n",
+      "  Chemo   | Subtype AUC=0.675 | TNM AUC=0.559 | DFS C-index=0.588 | OS C-index=0.476\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 002] Train Loss=4.770 | Val Loss=4.596\n",
+      "  All     | Subtype AUC=0.678 | TNM AUC=0.683 | DFS C-index=0.606 | OS C-index=0.536\n",
+      "  Immune  | Subtype AUC=0.691 | TNM AUC=0.654 | DFS C-index=0.585 | OS C-index=0.591\n",
+      "  Chemo   | Subtype AUC=0.634 | TNM AUC=0.729 | DFS C-index=0.639 | OS C-index=0.456\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 003] Train Loss=4.551 | Val Loss=4.404\n",
+      "  All     | Subtype AUC=0.698 | TNM AUC=0.678 | DFS C-index=0.592 | OS C-index=0.450\n",
+      "  Immune  | Subtype AUC=0.667 | TNM AUC=0.686 | DFS C-index=0.635 | OS C-index=0.480\n",
+      "  Chemo   | Subtype AUC=0.762 | TNM AUC=0.687 | DFS C-index=0.541 | OS C-index=0.379\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 004] Train Loss=4.299 | Val Loss=4.076\n",
+      "  All     | Subtype AUC=0.763 | TNM AUC=0.753 | DFS C-index=0.550 | OS C-index=0.597\n",
+      "  Immune  | Subtype AUC=0.720 | TNM AUC=0.745 | DFS C-index=0.554 | OS C-index=0.606\n",
+      "  Chemo   | Subtype AUC=0.824 | TNM AUC=0.735 | DFS C-index=0.510 | OS C-index=0.589\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 005] Train Loss=3.962 | Val Loss=3.941\n",
+      "  All     | Subtype AUC=0.693 | TNM AUC=0.703 | DFS C-index=0.640 | OS C-index=0.658\n",
+      "  Immune  | Subtype AUC=0.804 | TNM AUC=0.714 | DFS C-index=0.607 | OS C-index=0.675\n",
+      "  Chemo   | Subtype AUC=0.539 | TNM AUC=0.694 | DFS C-index=0.694 | OS C-index=0.621\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 006] Train Loss=3.911 | Val Loss=3.840\n",
+      "  All     | Subtype AUC=0.810 | TNM AUC=0.636 | DFS C-index=0.591 | OS C-index=0.549\n",
+      "  Immune  | Subtype AUC=0.840 | TNM AUC=0.657 | DFS C-index=0.552 | OS C-index=0.551\n",
+      "  Chemo   | Subtype AUC=0.771 | TNM AUC=0.618 | DFS C-index=0.580 | OS C-index=0.577\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 007] Train Loss=3.823 | Val Loss=3.876\n",
+      "  All     | Subtype AUC=0.800 | TNM AUC=0.722 | DFS C-index=0.645 | OS C-index=0.666\n",
+      "  Immune  | Subtype AUC=0.833 | TNM AUC=0.764 | DFS C-index=0.604 | OS C-index=0.654\n",
+      "  Chemo   | Subtype AUC=0.750 | TNM AUC=0.668 | DFS C-index=0.671 | OS C-index=0.605\n",
+      "\n",
+      "[Epoch 008] Train Loss=3.779 | Val Loss=3.794\n",
+      "  All     | Subtype AUC=0.689 | TNM AUC=0.711 | DFS C-index=0.654 | OS C-index=0.616\n",
+      "  Immune  | Subtype AUC=0.729 | TNM AUC=0.729 | DFS C-index=0.680 | OS C-index=0.682\n",
+      "  Chemo   | Subtype AUC=0.637 | TNM AUC=0.725 | DFS C-index=0.624 | OS C-index=0.565\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 009] Train Loss=3.725 | Val Loss=3.630\n",
+      "  All     | Subtype AUC=0.856 | TNM AUC=0.718 | DFS C-index=0.692 | OS C-index=0.677\n",
+      "  Immune  | Subtype AUC=0.897 | TNM AUC=0.791 | DFS C-index=0.713 | OS C-index=0.675\n",
+      "  Chemo   | Subtype AUC=0.785 | TNM AUC=0.652 | DFS C-index=0.678 | OS C-index=0.677\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 010] Train Loss=3.706 | Val Loss=3.646\n",
+      "  All     | Subtype AUC=0.749 | TNM AUC=0.773 | DFS C-index=0.659 | OS C-index=0.585\n",
+      "  Immune  | Subtype AUC=0.817 | TNM AUC=0.723 | DFS C-index=0.638 | OS C-index=0.596\n",
+      "  Chemo   | Subtype AUC=0.649 | TNM AUC=0.855 | DFS C-index=0.710 | OS C-index=0.540\n",
+      "\n",
+      "[Epoch 011] Train Loss=3.694 | Val Loss=3.690\n",
+      "  All     | Subtype AUC=0.791 | TNM AUC=0.769 | DFS C-index=0.637 | OS C-index=0.542\n",
+      "  Immune  | Subtype AUC=0.812 | TNM AUC=0.731 | DFS C-index=0.677 | OS C-index=0.593\n",
+      "  Chemo   | Subtype AUC=0.787 | TNM AUC=0.812 | DFS C-index=0.604 | OS C-index=0.488\n",
+      "\n",
+      "[Epoch 012] Train Loss=3.600 | Val Loss=3.469\n",
+      "  All     | Subtype AUC=0.875 | TNM AUC=0.744 | DFS C-index=0.663 | OS C-index=0.662\n",
+      "  Immune  | Subtype AUC=0.896 | TNM AUC=0.754 | DFS C-index=0.671 | OS C-index=0.753\n",
+      "  Chemo   | Subtype AUC=0.887 | TNM AUC=0.749 | DFS C-index=0.635 | OS C-index=0.573\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 013] Train Loss=3.641 | Val Loss=3.486\n",
+      "  All     | Subtype AUC=0.870 | TNM AUC=0.701 | DFS C-index=0.659 | OS C-index=0.649\n",
+      "  Immune  | Subtype AUC=0.912 | TNM AUC=0.701 | DFS C-index=0.708 | OS C-index=0.722\n",
+      "  Chemo   | Subtype AUC=0.811 | TNM AUC=0.724 | DFS C-index=0.600 | OS C-index=0.512\n",
+      "\n",
+      "[Epoch 014] Train Loss=3.554 | Val Loss=3.560\n",
+      "  All     | Subtype AUC=0.796 | TNM AUC=0.672 | DFS C-index=0.670 | OS C-index=0.656\n",
+      "  Immune  | Subtype AUC=0.773 | TNM AUC=0.739 | DFS C-index=0.685 | OS C-index=0.753\n",
+      "  Chemo   | Subtype AUC=0.818 | TNM AUC=0.600 | DFS C-index=0.659 | OS C-index=0.556\n",
+      "\n",
+      "[Epoch 015] Train Loss=3.490 | Val Loss=3.608\n",
+      "  All     | Subtype AUC=0.815 | TNM AUC=0.730 | DFS C-index=0.620 | OS C-index=0.622\n",
+      "  Immune  | Subtype AUC=0.850 | TNM AUC=0.846 | DFS C-index=0.657 | OS C-index=0.635\n",
+      "  Chemo   | Subtype AUC=0.752 | TNM AUC=0.590 | DFS C-index=0.565 | OS C-index=0.629\n",
+      "\n",
+      "[Epoch 016] Train Loss=3.559 | Val Loss=3.347\n",
+      "  All     | Subtype AUC=0.793 | TNM AUC=0.782 | DFS C-index=0.655 | OS C-index=0.636\n",
+      "  Immune  | Subtype AUC=0.889 | TNM AUC=0.781 | DFS C-index=0.643 | OS C-index=0.659\n",
+      "  Chemo   | Subtype AUC=0.673 | TNM AUC=0.809 | DFS C-index=0.663 | OS C-index=0.605\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 017] Train Loss=3.530 | Val Loss=3.405\n",
+      "  All     | Subtype AUC=0.840 | TNM AUC=0.731 | DFS C-index=0.645 | OS C-index=0.611\n",
+      "  Immune  | Subtype AUC=0.832 | TNM AUC=0.738 | DFS C-index=0.688 | OS C-index=0.654\n",
+      "  Chemo   | Subtype AUC=0.847 | TNM AUC=0.728 | DFS C-index=0.643 | OS C-index=0.597\n",
+      "\n",
+      "[Epoch 018] Train Loss=3.545 | Val Loss=3.434\n",
+      "  All     | Subtype AUC=0.912 | TNM AUC=0.729 | DFS C-index=0.660 | OS C-index=0.595\n",
+      "  Immune  | Subtype AUC=0.979 | TNM AUC=0.677 | DFS C-index=0.702 | OS C-index=0.585\n",
+      "  Chemo   | Subtype AUC=0.804 | TNM AUC=0.767 | DFS C-index=0.671 | OS C-index=0.605\n",
+      "\n",
+      "[Epoch 019] Train Loss=3.605 | Val Loss=3.394\n",
+      "  All     | Subtype AUC=0.807 | TNM AUC=0.722 | DFS C-index=0.655 | OS C-index=0.664\n",
+      "  Immune  | Subtype AUC=0.858 | TNM AUC=0.736 | DFS C-index=0.674 | OS C-index=0.701\n",
+      "  Chemo   | Subtype AUC=0.713 | TNM AUC=0.768 | DFS C-index=0.675 | OS C-index=0.645\n",
+      "\n",
+      "[Epoch 020] Train Loss=3.386 | Val Loss=3.381\n",
+      "  All     | Subtype AUC=0.784 | TNM AUC=0.757 | DFS C-index=0.674 | OS C-index=0.646\n",
+      "  Immune  | Subtype AUC=0.796 | TNM AUC=0.732 | DFS C-index=0.696 | OS C-index=0.727\n",
+      "  Chemo   | Subtype AUC=0.758 | TNM AUC=0.818 | DFS C-index=0.675 | OS C-index=0.565\n",
+      "\n",
+      "[Epoch 021] Train Loss=3.505 | Val Loss=3.388\n",
+      "  All     | Subtype AUC=0.767 | TNM AUC=0.836 | DFS C-index=0.679 | OS C-index=0.632\n",
+      "  Immune  | Subtype AUC=0.800 | TNM AUC=0.877 | DFS C-index=0.641 | OS C-index=0.696\n",
+      "  Chemo   | Subtype AUC=0.663 | TNM AUC=0.772 | DFS C-index=0.714 | OS C-index=0.544\n",
+      "\n",
+      "[Epoch 022] Train Loss=3.577 | Val Loss=3.618\n",
+      "  All     | Subtype AUC=0.723 | TNM AUC=0.683 | DFS C-index=0.641 | OS C-index=0.569\n",
+      "  Immune  | Subtype AUC=0.711 | TNM AUC=0.592 | DFS C-index=0.699 | OS C-index=0.619\n",
+      "  Chemo   | Subtype AUC=0.745 | TNM AUC=0.765 | DFS C-index=0.635 | OS C-index=0.524\n",
+      "\n",
+      "[Epoch 023] Train Loss=3.420 | Val Loss=3.244\n",
+      "  All     | Subtype AUC=0.874 | TNM AUC=0.752 | DFS C-index=0.654 | OS C-index=0.653\n",
+      "  Immune  | Subtype AUC=0.912 | TNM AUC=0.839 | DFS C-index=0.652 | OS C-index=0.701\n",
+      "  Chemo   | Subtype AUC=0.833 | TNM AUC=0.629 | DFS C-index=0.639 | OS C-index=0.625\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 024] Train Loss=3.525 | Val Loss=3.595\n",
+      "  All     | Subtype AUC=0.775 | TNM AUC=0.754 | DFS C-index=0.600 | OS C-index=0.610\n",
+      "  Immune  | Subtype AUC=0.733 | TNM AUC=0.758 | DFS C-index=0.641 | OS C-index=0.690\n",
+      "  Chemo   | Subtype AUC=0.800 | TNM AUC=0.773 | DFS C-index=0.631 | OS C-index=0.589\n",
+      "\n",
+      "[Epoch 025] Train Loss=3.571 | Val Loss=3.480\n",
+      "  All     | Subtype AUC=0.820 | TNM AUC=0.763 | DFS C-index=0.599 | OS C-index=0.627\n",
+      "  Immune  | Subtype AUC=0.825 | TNM AUC=0.742 | DFS C-index=0.557 | OS C-index=0.627\n",
+      "  Chemo   | Subtype AUC=0.837 | TNM AUC=0.796 | DFS C-index=0.616 | OS C-index=0.613\n",
+      "\n",
+      "[Epoch 026] Train Loss=3.508 | Val Loss=3.209\n",
+      "  All     | Subtype AUC=0.874 | TNM AUC=0.765 | DFS C-index=0.724 | OS C-index=0.633\n",
+      "  Immune  | Subtype AUC=0.853 | TNM AUC=0.804 | DFS C-index=0.713 | OS C-index=0.688\n",
+      "  Chemo   | Subtype AUC=0.903 | TNM AUC=0.719 | DFS C-index=0.788 | OS C-index=0.573\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 027] Train Loss=3.417 | Val Loss=3.468\n",
+      "  All     | Subtype AUC=0.836 | TNM AUC=0.674 | DFS C-index=0.685 | OS C-index=0.672\n",
+      "  Immune  | Subtype AUC=0.834 | TNM AUC=0.707 | DFS C-index=0.713 | OS C-index=0.701\n",
+      "  Chemo   | Subtype AUC=0.873 | TNM AUC=0.604 | DFS C-index=0.608 | OS C-index=0.589\n",
+      "\n",
+      "[Epoch 028] Train Loss=3.443 | Val Loss=3.314\n",
+      "  All     | Subtype AUC=0.817 | TNM AUC=0.734 | DFS C-index=0.653 | OS C-index=0.670\n",
+      "  Immune  | Subtype AUC=0.875 | TNM AUC=0.738 | DFS C-index=0.604 | OS C-index=0.719\n",
+      "  Chemo   | Subtype AUC=0.744 | TNM AUC=0.745 | DFS C-index=0.694 | OS C-index=0.593\n",
+      "\n",
+      "[Epoch 029] Train Loss=3.451 | Val Loss=3.259\n",
+      "  All     | Subtype AUC=0.782 | TNM AUC=0.819 | DFS C-index=0.708 | OS C-index=0.637\n",
+      "  Immune  | Subtype AUC=0.814 | TNM AUC=0.811 | DFS C-index=0.638 | OS C-index=0.688\n",
+      "  Chemo   | Subtype AUC=0.750 | TNM AUC=0.833 | DFS C-index=0.776 | OS C-index=0.552\n",
+      "\n",
+      "[Epoch 030] Train Loss=3.556 | Val Loss=3.322\n",
+      "  All     | Subtype AUC=0.880 | TNM AUC=0.762 | DFS C-index=0.650 | OS C-index=0.662\n",
+      "  Immune  | Subtype AUC=0.884 | TNM AUC=0.756 | DFS C-index=0.660 | OS C-index=0.727\n",
+      "  Chemo   | Subtype AUC=0.863 | TNM AUC=0.773 | DFS C-index=0.584 | OS C-index=0.548\n",
+      "\n",
+      "[Epoch 031] Train Loss=3.493 | Val Loss=3.233\n",
+      "  All     | Subtype AUC=0.857 | TNM AUC=0.764 | DFS C-index=0.705 | OS C-index=0.638\n",
+      "  Immune  | Subtype AUC=0.894 | TNM AUC=0.772 | DFS C-index=0.663 | OS C-index=0.688\n",
+      "  Chemo   | Subtype AUC=0.800 | TNM AUC=0.755 | DFS C-index=0.757 | OS C-index=0.585\n",
+      "\n",
+      "[Epoch 032] Train Loss=3.483 | Val Loss=3.330\n",
+      "  All     | Subtype AUC=0.894 | TNM AUC=0.743 | DFS C-index=0.616 | OS C-index=0.595\n",
+      "  Immune  | Subtype AUC=0.917 | TNM AUC=0.792 | DFS C-index=0.613 | OS C-index=0.567\n",
+      "  Chemo   | Subtype AUC=0.869 | TNM AUC=0.644 | DFS C-index=0.561 | OS C-index=0.597\n",
+      "\n",
+      "[Epoch 033] Train Loss=3.447 | Val Loss=3.307\n",
+      "  All     | Subtype AUC=0.824 | TNM AUC=0.846 | DFS C-index=0.657 | OS C-index=0.660\n",
+      "  Immune  | Subtype AUC=0.936 | TNM AUC=0.821 | DFS C-index=0.635 | OS C-index=0.740\n",
+      "  Chemo   | Subtype AUC=0.655 | TNM AUC=0.898 | DFS C-index=0.671 | OS C-index=0.577\n",
+      "\n",
+      "[Epoch 034] Train Loss=3.396 | Val Loss=3.207\n",
+      "  All     | Subtype AUC=0.779 | TNM AUC=0.746 | DFS C-index=0.679 | OS C-index=0.693\n",
+      "  Immune  | Subtype AUC=0.747 | TNM AUC=0.774 | DFS C-index=0.641 | OS C-index=0.719\n",
+      "  Chemo   | Subtype AUC=0.855 | TNM AUC=0.703 | DFS C-index=0.584 | OS C-index=0.605\n",
+      "  ✓ Best model updated\n",
+      "\n",
+      "[Epoch 035] Train Loss=3.362 | Val Loss=3.279\n",
+      "  All     | Subtype AUC=0.802 | TNM AUC=0.756 | DFS C-index=0.666 | OS C-index=0.667\n",
+      "  Immune  | Subtype AUC=0.861 | TNM AUC=0.826 | DFS C-index=0.749 | OS C-index=0.709\n",
+      "  Chemo   | Subtype AUC=0.733 | TNM AUC=0.658 | DFS C-index=0.580 | OS C-index=0.609\n",
+      "\n",
+      "[Epoch 036] Train Loss=3.468 | Val Loss=3.541\n",
+      "  All     | Subtype AUC=0.737 | TNM AUC=0.736 | DFS C-index=0.614 | OS C-index=0.601\n",
+      "  Immune  | Subtype AUC=0.764 | TNM AUC=0.743 | DFS C-index=0.579 | OS C-index=0.635\n",
+      "  Chemo   | Subtype AUC=0.725 | TNM AUC=0.736 | DFS C-index=0.616 | OS C-index=0.565\n",
+      "\n",
+      "[Epoch 037] Train Loss=3.259 | Val Loss=3.274\n",
+      "  All     | Subtype AUC=0.883 | TNM AUC=0.793 | DFS C-index=0.654 | OS C-index=0.630\n",
+      "  Immune  | Subtype AUC=0.857 | TNM AUC=0.751 | DFS C-index=0.660 | OS C-index=0.675\n",
+      "  Chemo   | Subtype AUC=0.905 | TNM AUC=0.837 | DFS C-index=0.655 | OS C-index=0.569\n",
+      "\n",
+      "[Epoch 038] Train Loss=3.328 | Val Loss=3.333\n",
+      "  All     | Subtype AUC=0.809 | TNM AUC=0.772 | DFS C-index=0.655 | OS C-index=0.630\n",
+      "  Immune  | Subtype AUC=0.861 | TNM AUC=0.813 | DFS C-index=0.652 | OS C-index=0.646\n",
+      "  Chemo   | Subtype AUC=0.758 | TNM AUC=0.719 | DFS C-index=0.647 | OS C-index=0.589\n",
+      "\n",
+      "[Epoch 039] Train Loss=3.409 | Val Loss=3.216\n",
+      "  All     | Subtype AUC=0.871 | TNM AUC=0.768 | DFS C-index=0.679 | OS C-index=0.625\n",
+      "  Immune  | Subtype AUC=0.882 | TNM AUC=0.744 | DFS C-index=0.783 | OS C-index=0.703\n",
+      "  Chemo   | Subtype AUC=0.865 | TNM AUC=0.823 | DFS C-index=0.624 | OS C-index=0.544\n",
+      "\n",
+      "[Epoch 040] Train Loss=3.421 | Val Loss=3.211\n",
+      "  All     | Subtype AUC=0.883 | TNM AUC=0.694 | DFS C-index=0.737 | OS C-index=0.680\n",
+      "  Immune  | Subtype AUC=0.901 | TNM AUC=0.700 | DFS C-index=0.674 | OS C-index=0.701\n",
+      "  Chemo   | Subtype AUC=0.870 | TNM AUC=0.688 | DFS C-index=0.741 | OS C-index=0.573\n",
+      "\n",
+      "[Epoch 041] Train Loss=3.318 | Val Loss=3.326\n",
+      "  All     | Subtype AUC=0.821 | TNM AUC=0.778 | DFS C-index=0.680 | OS C-index=0.678\n",
+      "  Immune  | Subtype AUC=0.855 | TNM AUC=0.817 | DFS C-index=0.613 | OS C-index=0.722\n",
+      "  Chemo   | Subtype AUC=0.781 | TNM AUC=0.724 | DFS C-index=0.769 | OS C-index=0.661\n",
+      "\n",
+      "[Epoch 042] Train Loss=3.203 | Val Loss=3.307\n",
+      "  All     | Subtype AUC=0.825 | TNM AUC=0.778 | DFS C-index=0.677 | OS C-index=0.603\n",
+      "  Immune  | Subtype AUC=0.851 | TNM AUC=0.770 | DFS C-index=0.708 | OS C-index=0.609\n",
+      "  Chemo   | Subtype AUC=0.850 | TNM AUC=0.815 | DFS C-index=0.675 | OS C-index=0.597\n",
+      "\n",
+      "⏹ Early stopping triggered\n",
+      "\n",
+      "Running inference with best model...\n",
+      "train | loss=3.330 | Immune=105 Chemo=75\n",
+      "val   | loss=3.458 | Immune=34 Chemo=26\n",
+      "test  | loss=3.877 | Immune=32 Chemo=28\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "findfont: Font family ['Arial'] not found. Falling back to DejaVu Sans.\n",
+      "findfont: Font family ['Arial'] not found. Falling back to DejaVu Sans.\n",
+      "findfont: Font family ['Arial'] not found. Falling back to DejaVu Sans.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✔ Figure 7 generated (DFS/OS KM + HR) for Immune/Chemo.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0, \"/export/home/daifang/lunghospital/MM-DLS-master/MM-DLS-master\")\n",
+    "# main.py\n",
+    "import os\n",
+    "import sys\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.utils.data import DataLoader, random_split\n",
+    "\n",
+    "from sklearn.metrics import roc_auc_score, accuracy_score\n",
+    "from sklearn.preprocessing import label_binarize\n",
+    "\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "from lifelines import KaplanMeierFitter, CoxPHFitter\n",
+    "from lifelines.statistics import multivariate_logrank_test\n",
+    "from lifelines.utils import concordance_index\n",
+    "from sklearn.metrics import brier_score_loss\n",
+    "from scipy.stats import norm\n",
+    "\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# Project path (IMPORTANT for Jupyter / HPC)\n",
+    "# =========================================================\n",
+    "PROJECT_ROOT = os.path.abspath(\".\")\n",
+    "if PROJECT_ROOT not in sys.path:\n",
+    "    sys.path.insert(0, PROJECT_ROOT)\n",
+    "\n",
+    "# =========================================================\n",
+    "# imports:  mm_dls/ \n",
+    "# =========================================================\n",
+    "def _import_modules():\n",
+    "\n",
+    "    from mm_dls.HierMM_DLS import HierMM_DLS\n",
+    "    from mm_dls.FakePatientDataset import FakePatientDataset\n",
+    "    from mm_dls.CoxphLoss import CoxPHLoss\n",
+    "    return HierMM_DLS, FakePatientDataset, CoxPHLoss\n",
+    "\n",
+    "\n",
+    "HierMM_DLS, FakePatientDataset, CoxPHLoss = _import_modules()\n",
+    "\n",
+    "\n",
+    "# =========================\n",
+    "# Training configuration\n",
+    "# =========================\n",
+    "EPOCHS        = 300\n",
+    "PATIENCE      = 8\n",
+    "BATCH_SIZE    = 4\n",
+    "LR            = 1e-4\n",
+    "WEIGHT_DECAY  = 1e-5\n",
+    "\n",
+    "# =========================\n",
+    "# Task definition\n",
+    "# =========================\n",
+    "NUM_SUBTYPES  = 2        # e.g., LUAD vs LUSC\n",
+    "NUM_TNM       = 3        # Stage I–II / III / IV\n",
+    "\n",
+    "# =========================\n",
+    "# Image settings\n",
+    "# =========================\n",
+    "N_SLICES      = 30       # max slices per patient\n",
+    "IMG_SIZE      = 224\n",
+    "\n",
+    "\n",
+    "SAVE_DIR = \"./results\"\n",
+    "FIG_DIR  = \"./figures\"\n",
+    "os.makedirs(SAVE_DIR, exist_ok=True)\n",
+    "os.makedirs(FIG_DIR, exist_ok=True)\n",
+    "\n",
+    "# -------------------------\n",
+    "# GPU (force cuda:1)\n",
+    "# -------------------------\n",
+    "assert torch.cuda.is_available(), \"CUDA not available\"\n",
+    "DEVICE = torch.device(\"cuda:1\")\n",
+    "torch.cuda.set_device(DEVICE)\n",
+    "print(\"Using device:\", DEVICE)\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# Core utils\n",
+    "# =========================================================\n",
+    "def _sigmoid(x):\n",
+    "    return 1 / (1 + np.exp(-x))\n",
+    "\n",
+    "def _ensure_numpy(x):\n",
+    "    if isinstance(x, torch.Tensor):\n",
+    "        return x.detach().cpu().numpy()\n",
+    "    return x\n",
+    "\n",
+    "def _risk_to_groups(risk, q=(1/3, 2/3), labels=(\"Low\", \"Mediate\", \"High\")):\n",
+    "    \"\"\"\n",
+    "    Convert continuous risk into 3 groups by tertiles.\n",
+    "    \"\"\"\n",
+    "    r = np.asarray(risk).reshape(-1)\n",
+    "    t1, t2 = np.quantile(r, q[0]), np.quantile(r, q[1])\n",
+    "    out = np.full(len(r), labels[1], dtype=object)\n",
+    "    out[r <= t1] = labels[0]\n",
+    "    out[r >= t2] = labels[2]\n",
+    "    return out\n",
+    "\n",
+    "def _evaluate_survival_metrics(time, event, risk, time_point=30):\n",
+    "    \"\"\"\n",
+    "    C-index + Brier at a fixed time point.\n",
+    "    risk: higher => earlier event, so use -risk in concordance_index.\n",
+    "    \"\"\"\n",
+    "    time = np.asarray(time).reshape(-1)\n",
+    "    event = np.asarray(event).reshape(-1).astype(int)\n",
+    "    risk = np.asarray(risk).reshape(-1)\n",
+    "\n",
+    "    c_index = concordance_index(time, -risk, event)\n",
+    "\n",
+    "    # Brier: predict survival at time_point using a monotonic transform of risk (proxy)\n",
+    "    # This is a \"proxy\" survival probability for demo/debug; replace with proper survival model if needed.\n",
+    "    y_true = (time > time_point).astype(int)  # 1 means survived beyond time_point\n",
+    "    # map risk into [0,1] survival prob proxy: higher risk => lower survival prob\n",
+    "    y_prob = 1 - (risk - risk.min()) / (risk.max() - risk.min() + 1e-8)\n",
+    "    brier = brier_score_loss(y_true, y_prob)\n",
+    "\n",
+    "    return float(c_index), float(brier)\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# One epoch (train / eval)\n",
+    "# =========================================================\n",
+    "def run_epoch_verbose(model, loader, optimizer, device, train=True):\n",
+    "    ce  = nn.CrossEntropyLoss()\n",
+    "    bce = nn.BCEWithLogitsLoss(reduction=\"none\")\n",
+    "    cox = CoxPHLoss()\n",
+    "\n",
+    "    model.train() if train else model.eval()\n",
+    "\n",
+    "    losses = []\n",
+    "\n",
+    "    # classification\n",
+    "    sub_y_all, sub_s_all = [], []\n",
+    "    tnm_y_all, tnm_s_all = [], []\n",
+    "    treat_all = []\n",
+    "\n",
+    "    # survival (cox risk + time/event)\n",
+    "    dfs_r_all, dfs_t_all, dfs_e_all = [], [], []\n",
+    "    os_r_all,  os_t_all,  os_e_all  = [], [], []\n",
+    "\n",
+    "    # survival 1y/3y/5y logits (optional save)\n",
+    "    dfs_log_all, os_log_all = [], []\n",
+    "\n",
+    "    for batch in loader:\n",
+    "        # NOTE: dataset must return 19 items including treatment\n",
+    "        if len(batch) != 19:\n",
+    "            raise ValueError(f\"Batch length mismatch: expected 19, got {len(batch)}. \"\n",
+    "                             f\"Please ensure Dataset __getitem__ returns treatment as the 19th item.\")\n",
+    "\n",
+    "        (\n",
+    "            pid, lesion, space, rad, pet, cli,\n",
+    "            y_sub, y_tnm,\n",
+    "            dfs_t, dfs_e,\n",
+    "            os_t, os_e,\n",
+    "            dfs1, dfs3, dfs5,\n",
+    "            os1, os3, os5,\n",
+    "            treatment\n",
+    "        ) = batch\n",
+    "\n",
+    "        lesion, space = lesion.to(device), space.to(device)\n",
+    "        rad, pet, cli = rad.to(device), pet.to(device), cli.to(device)\n",
+    "        y_sub, y_tnm  = y_sub.to(device), y_tnm.to(device)\n",
+    "        dfs_t, dfs_e = dfs_t.to(device), dfs_e.to(device)\n",
+    "        os_t,  os_e  = os_t.to(device),  os_e.to(device)\n",
+    "        treatment = treatment.to(device)\n",
+    "\n",
+    "        dfs_y = torch.stack([dfs1, dfs3, dfs5], dim=1).to(device)\n",
+    "        os_y  = torch.stack([os1,  os3,  os5 ], dim=1).to(device)\n",
+    "\n",
+    "        with torch.set_grad_enabled(train):\n",
+    "            sub_l, tnm_l, dfs_r, os_r, dfs_log, os_log = model(\n",
+    "                lesion, space, rad, pet, cli\n",
+    "            )\n",
+    "\n",
+    "            loss = (\n",
+    "                ce(sub_l, y_sub) +\n",
+    "                ce(tnm_l, y_tnm) +\n",
+    "                cox(dfs_r, dfs_t, dfs_e) +\n",
+    "                cox(os_r,  os_t,  os_e) +\n",
+    "                bce(dfs_log, dfs_y).mean() +\n",
+    "                bce(os_log,  os_y ).mean()\n",
+    "            )\n",
+    "\n",
+    "            if train:\n",
+    "                optimizer.zero_grad()\n",
+    "                loss.backward()\n",
+    "                optimizer.step()\n",
+    "\n",
+    "        losses.append(loss.item())\n",
+    "\n",
+    "        # ----- Collect predictions -----\n",
+    "        sub_prob = torch.softmax(sub_l, dim=1)[:, 1]     # subtype prob\n",
+    "        tnm_prob = torch.softmax(tnm_l, dim=1)           # [B,3]\n",
+    "\n",
+    "        sub_s_all.append(_ensure_numpy(sub_prob))\n",
+    "        sub_y_all.append(_ensure_numpy(y_sub))\n",
+    "\n",
+    "        tnm_s_all.append(_ensure_numpy(tnm_prob))\n",
+    "        tnm_y_all.append(_ensure_numpy(y_tnm))\n",
+    "\n",
+    "        treat_all.append(_ensure_numpy(treatment))\n",
+    "\n",
+    "        # survival\n",
+    "        dfs_r_all.append(_ensure_numpy(dfs_r))\n",
+    "        dfs_t_all.append(_ensure_numpy(dfs_t))\n",
+    "        dfs_e_all.append(_ensure_numpy(dfs_e))\n",
+    "\n",
+    "        os_r_all.append(_ensure_numpy(os_r))\n",
+    "        os_t_all.append(_ensure_numpy(os_t))\n",
+    "        os_e_all.append(_ensure_numpy(os_e))\n",
+    "\n",
+    "        dfs_log_all.append(_ensure_numpy(dfs_log))\n",
+    "        os_log_all.append(_ensure_numpy(os_log))\n",
+    "\n",
+    "    return (\n",
+    "        float(np.mean(losses)),\n",
+    "\n",
+    "        np.concatenate(sub_y_all),\n",
+    "        np.concatenate(sub_s_all),\n",
+    "\n",
+    "        np.concatenate(tnm_y_all),\n",
+    "        np.concatenate(tnm_s_all),\n",
+    "\n",
+    "        np.concatenate(treat_all),\n",
+    "\n",
+    "        np.concatenate(dfs_r_all),\n",
+    "        np.concatenate(dfs_t_all),\n",
+    "        np.concatenate(dfs_e_all),\n",
+    "\n",
+    "        np.concatenate(os_r_all),\n",
+    "        np.concatenate(os_t_all),\n",
+    "        np.concatenate(os_e_all),\n",
+    "\n",
+    "        np.concatenate(dfs_log_all, axis=0),  # [N,3]\n",
+    "        np.concatenate(os_log_all, axis=0),   # [N,3]\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# Evaluation by cohort (classification + survival)\n",
+    "# =========================================================\n",
+    "def evaluate_by_treatment(sub_y, sub_s, tnm_y, tnm_s, treat,\n",
+    "                          dfs_r, dfs_t, dfs_e, os_r, os_t, os_e):\n",
+    "    results = {}\n",
+    "\n",
+    "    cohorts = {\n",
+    "        \"All\": np.ones_like(treat, dtype=bool),\n",
+    "        \"Immune\": treat == 0,\n",
+    "        \"Chemo\":  treat == 1,\n",
+    "    }\n",
+    "\n",
+    "    for name, mask in cohorts.items():\n",
+    "        if mask.sum() < 10:\n",
+    "            continue\n",
+    "\n",
+    "        res = {}\n",
+    "\n",
+    "        # Subtype (binary)\n",
+    "        res[\"Subtype_AUC\"] = roc_auc_score(sub_y[mask], sub_s[mask])\n",
+    "        res[\"Subtype_ACC\"] = accuracy_score(sub_y[mask], (sub_s[mask] > 0.5).astype(int))\n",
+    "\n",
+    "        # TNM (multiclass macro AUC + ACC)\n",
+    "        tnm_bin = label_binarize(tnm_y[mask], classes=[0, 1, 2])\n",
+    "        res[\"TNM_AUC_macro\"] = roc_auc_score(\n",
+    "            tnm_bin, tnm_s[mask], average=\"macro\", multi_class=\"ovr\"\n",
+    "        )\n",
+    "        res[\"TNM_ACC\"] = accuracy_score(\n",
+    "            tnm_y[mask], np.argmax(tnm_s[mask], axis=1)\n",
+    "        )\n",
+    "\n",
+    "        # Survival\n",
+    "        dfs_c, dfs_b = _evaluate_survival_metrics(dfs_t[mask], dfs_e[mask], dfs_r[mask], time_point=30)\n",
+    "        os_c,  os_b  = _evaluate_survival_metrics(os_t[mask],  os_e[mask],  os_r[mask],  time_point=30)\n",
+    "\n",
+    "        res[\"DFS_C_index\"] = dfs_c\n",
+    "        res[\"DFS_Brier_30m\"] = dfs_b\n",
+    "        res[\"OS_C_index\"] = os_c\n",
+    "        res[\"OS_Brier_30m\"] = os_b\n",
+    "\n",
+    "        results[name] = res\n",
+    "\n",
+    "    return results\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# Figure 7: KM + HR (per cohort, per endpoint)\n",
+    "# =========================================================\n",
+    "def plot_km_curve_with_hr(df, title, save_prefix):\n",
+    "    \"\"\"\n",
+    "    df must contain columns: time, event, group (Low/Mediate/High)\n",
+    "    \"\"\"\n",
+    "    kmf = KaplanMeierFitter()\n",
+    "    fig, ax = plt.subplots(figsize=(8, 6), facecolor=\"white\")\n",
+    "    ax.set_facecolor(\"white\")\n",
+    "\n",
+    "    colors = {\"Low\": \"#91c7ae\", \"Mediate\": \"#f7b977\", \"High\": \"#d87c7c\"}\n",
+    "    groups = [\"Low\", \"Mediate\", \"High\"]\n",
+    "\n",
+    "    # plot KM\n",
+    "    lines = {}\n",
+    "    at_risk_table = []\n",
+    "    times = np.arange(0, 70, 10)\n",
+    "\n",
+    "    for g in groups:\n",
+    "        m = df[\"group\"] == g\n",
+    "        if m.sum() == 0:\n",
+    "            continue\n",
+    "\n",
+    "        kmf.fit(df.loc[m, \"time\"], event_observed=df.loc[m, \"event\"], label=g)\n",
+    "        kmf.plot_survival_function(\n",
+    "            ax=ax, ci_show=True, linewidth=2, color=colors[g], marker=\"+\"\n",
+    "        )\n",
+    "        lines[g] = ax.get_lines()[-1]\n",
+    "\n",
+    "        at_risk_table.append([np.sum(df.loc[m, \"time\"] >= t) for t in times])\n",
+    "\n",
+    "    # legend\n",
+    "    handles = [lines[g] for g in groups if g in lines]\n",
+    "    labels = [\"Low\", \"Medium\", \"High\"][:len(handles)]\n",
+    "    ax.legend(handles, labels, title=\"Groups\", loc=\"upper right\",\n",
+    "              frameon=True, framealpha=0.5, fontsize=12, title_fontsize=12)\n",
+    "\n",
+    "    # at risk numbers (optional, matches your style)\n",
+    "    if len(at_risk_table) == 3:\n",
+    "        low, mid, high = at_risk_table\n",
+    "        for i, t in enumerate(times):\n",
+    "            ax.text(t, -0.38, str(low[i]),  color=\"#207f4c\", fontsize=14, ha=\"center\")\n",
+    "            ax.text(t, -0.48, str(mid[i]),  color=\"#fca106\", fontsize=14, ha=\"center\")\n",
+    "            ax.text(t, -0.58, str(high[i]), color=\"#cc163a\", fontsize=14, ha=\"center\")\n",
+    "\n",
+    "        ax.text(-1,  -0.28, \"Number at risk\", color=\"black\", ha=\"center\", fontsize=14)\n",
+    "        ax.text(-10, -0.38, \"Low\",    color=\"#207f4c\", fontsize=14)\n",
+    "        ax.text(-10, -0.48, \"Medium\", color=\"#fca106\", fontsize=14)\n",
+    "        ax.text(-10, -0.58, \"High\",   color=\"#cc163a\", fontsize=14)\n",
+    "\n",
+    "    # Cox HR + p-values\n",
+    "    df2 = df.copy()\n",
+    "    df2[\"group_code\"] = df2[\"group\"].map({\"Low\": 0, \"Mediate\": 1, \"High\": 2})\n",
+    "    cph = CoxPHFitter()\n",
+    "    cph.fit(df2[[\"time\", \"event\", \"group_code\"]], duration_col=\"time\", event_col=\"event\")\n",
+    "\n",
+    "    coef = float(cph.params_[\"group_code\"])\n",
+    "    se   = float(cph.standard_errors_[\"group_code\"])\n",
+    "\n",
+    "    hr_med_vs_low  = np.exp(coef * 1)\n",
+    "    hr_high_vs_low = np.exp(coef * 2)\n",
+    "\n",
+    "    z_med  = (coef * 1) / se\n",
+    "    p_med  = 2 * (1 - norm.cdf(abs(z_med)))\n",
+    "\n",
+    "    z_high = (coef * 2) / se\n",
+    "    p_high = 2 * (1 - norm.cdf(abs(z_high)))\n",
+    "\n",
+    "    # logrank\n",
+    "    res_lr = multivariate_logrank_test(df2[\"time\"], df2[\"group\"], df2[\"event\"])\n",
+    "\n",
+    "    # C-index + brier (proxy)\n",
+    "    c_index, brier = _evaluate_survival_metrics(df2[\"time\"].values, df2[\"event\"].values,\n",
+    "                                                df2[\"group_code\"].values, time_point=30)\n",
+    "\n",
+    "    ax.text(25, 0.46, f\"P(log-rank)={res_lr.p_value:.3f}\", fontsize=12)\n",
+    "    ax.text(25, 0.36, f\"C-index={c_index:.3f}\", fontsize=12)\n",
+    "    ax.text(25, 0.26, f\"Brier(30m)={brier:.3f}\", fontsize=12)\n",
+    "    ax.text(25, 0.16, f\"HR Intermediate vs Low = {hr_med_vs_low:.2f}, P={p_med:.3f}\", fontsize=12)\n",
+    "    ax.text(25, 0.06, f\"HR High vs Low = {hr_high_vs_low:.2f}, P={p_high:.3f}\", fontsize=12)\n",
+    "\n",
+    "    # cosmetics\n",
+    "    ax.spines[\"top\"].set_visible(False)\n",
+    "    ax.spines[\"right\"].set_visible(False)\n",
+    "    ax.set_title(title, fontsize=14)\n",
+    "    ax.set_xlabel(\"Time since treatment start (months)\", fontsize=14)\n",
+    "    ax.set_ylabel(\"Survival probability\", fontsize=14)\n",
+    "    ax.set_ylim(0, 1.05)\n",
+    "    ax.grid(alpha=0.3)\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(save_prefix + \".png\", dpi=600, bbox_inches=\"tight\")\n",
+    "    plt.savefig(save_prefix + \".pdf\", dpi=600, bbox_inches=\"tight\")\n",
+    "    plt.close()\n",
+    "    return save_prefix\n",
+    "\n",
+    "\n",
+    "def generate_figure_from_saved(result_dir=SAVE_DIR, fig_dir=FIG_DIR, which_split=(\"val\", \"test\")):\n",
+    "    \"\"\"\n",
+    "    Load saved dfs/os arrays and generate KM+HR for Immune/Chemo separately.\n",
+    "    \"\"\"\n",
+    "    os.makedirs(fig_dir, exist_ok=True)\n",
+    "\n",
+    "    for split in which_split:\n",
+    "        # load arrays\n",
+    "        trt = np.load(os.path.join(result_dir, f\"treatment_{split}.npy\"))\n",
+    "\n",
+    "        dfs_r = np.load(os.path.join(result_dir, f\"dfs_{split}_risk.npy\"))\n",
+    "        dfs_t = np.load(os.path.join(result_dir, f\"dfs_{split}_time.npy\"))\n",
+    "        dfs_e = np.load(os.path.join(result_dir, f\"dfs_{split}_event.npy\"))\n",
+    "\n",
+    "        os_r  = np.load(os.path.join(result_dir, f\"os_{split}_risk.npy\"))\n",
+    "        os_t  = np.load(os.path.join(result_dir, f\"os_{split}_time.npy\"))\n",
+    "        os_e  = np.load(os.path.join(result_dir, f\"os_{split}_event.npy\"))\n",
+    "\n",
+    "        for cohort_name, mask in {\n",
+    "            \"Immune\": trt == 0,\n",
+    "            \"Chemo\":  trt == 1\n",
+    "        }.items():\n",
+    "            if mask.sum() < 20:\n",
+    "                print(f\"[Figure7] Skip {split}-{cohort_name}: too few samples ({mask.sum()})\")\n",
+    "                continue\n",
+    "\n",
+    "            # DFS groups\n",
+    "            dfs_group = _risk_to_groups(dfs_r[mask])\n",
+    "            df_dfs = pd.DataFrame({\n",
+    "                \"time\": dfs_t[mask],\n",
+    "                \"event\": dfs_e[mask].astype(int),\n",
+    "                \"group\": dfs_group\n",
+    "            })\n",
+    "\n",
+    "            # OS groups\n",
+    "            os_group = _risk_to_groups(os_r[mask])\n",
+    "            df_os = pd.DataFrame({\n",
+    "                \"time\": os_t[mask],\n",
+    "                \"event\": os_e[mask].astype(int),\n",
+    "                \"group\": os_group\n",
+    "            })\n",
+    "\n",
+    "            # save CSV (optional, for reproducibility)\n",
+    "            df_dfs.to_csv(os.path.join(result_dir, f\"dfs_{split}_{cohort_name}.csv\"), index=False)\n",
+    "            df_os.to_csv(os.path.join(result_dir, f\"os_{split}_{cohort_name}.csv\"), index=False)\n",
+    "\n",
+    "            # plot\n",
+    "            plot_km_curve_with_hr(\n",
+    "                df_dfs,\n",
+    "                title=f\"Disease-Free Survival (DFS) — Kaplan-Meier Curves\\n{cohort_name} {split} set (n={mask.sum()})\",\n",
+    "                save_prefix=os.path.join(fig_dir, f\"Figure7_DFS_{cohort_name}_{split}\")\n",
+    "            )\n",
+    "            plot_km_curve_with_hr(\n",
+    "                df_os,\n",
+    "                title=f\"Overall Survival (OS) — Kaplan-Meier Curves\\n{cohort_name} {split} set (n={mask.sum()})\",\n",
+    "                save_prefix=os.path.join(fig_dir, f\"Figure7_OS_{cohort_name}_{split}\")\n",
+    "            )\n",
+    "\n",
+    "    print(\"✔ Figure 7 generated (DFS/OS KM + HR) for Immune/Chemo.\")\n",
+    "\n",
+    "\n",
+    "# =========================================================\n",
+    "# Main\n",
+    "# =========================================================\n",
+    "def main():\n",
+    "    # -------------------------\n",
+    "    # Dataset (must return treatment as 19th item)\n",
+    "    # -------------------------\n",
+    "    from mm_dls.PatientDataset import PatientDataset\n",
+    "\n",
+    "    dataset = PatientDataset(\n",
+    "        data_root=\"/path/to/DATA_ROOT\",\n",
+    "        clinical_csv=\"/path/to/clinical.csv\",\n",
+    "        radiomics_npy=\"/path/to/radiomics.npy\",\n",
+    "        pet_npy=\"/path/to/pet.npy\",\n",
+    "        n_slices=N_SLICES,\n",
+    "        img_size=IMG_SIZE,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "    n_train = int(0.6 * len(dataset))\n",
+    "    n_val   = int(0.2 * len(dataset))\n",
+    "    n_test  = len(dataset) - n_train - n_val\n",
+    "\n",
+    "    train_set, val_set, test_set = random_split(dataset, [n_train, n_val, n_test])\n",
+    "\n",
+    "    loaders = {\n",
+    "        \"train\": DataLoader(train_set, BATCH_SIZE, shuffle=True,  num_workers=4),\n",
+    "        \"val\":   DataLoader(val_set,   BATCH_SIZE, shuffle=False, num_workers=4),\n",
+    "        \"test\":  DataLoader(test_set,  BATCH_SIZE, shuffle=False, num_workers=4),\n",
+    "    }\n",
+    "\n",
+    "    # -------------------------\n",
+    "    # Model\n",
+    "    # -------------------------\n",
+    "    model = HierMM_DLS(NUM_SUBTYPES, NUM_TNM).to(DEVICE)\n",
+    "    optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n",
+    "\n",
+    "    best_val_loss = 1e9\n",
+    "    wait = 0\n",
+    "\n",
+    "    # -------------------------\n",
+    "    # Training\n",
+    "    # -------------------------\n",
+    "    for epoch in range(1, EPOCHS + 1):\n",
+    "        tr = run_epoch_verbose(model, loaders[\"train\"], optimizer, DEVICE, train=True)\n",
+    "        va = run_epoch_verbose(model, loaders[\"val\"],   optimizer, DEVICE, train=False)\n",
+    "\n",
+    "        tr_loss = tr[0]\n",
+    "        va_loss = va[0]\n",
+    "\n",
+    "        # unpack val for metrics\n",
+    "        _, sy, ss, ty, ts, trt, dfs_r, dfs_t, dfs_e, os_r, os_t, os_e, _, _ = va\n",
+    "        metrics = evaluate_by_treatment(sy, ss, ty, ts, trt, dfs_r, dfs_t, dfs_e, os_r, os_t, os_e)\n",
+    "\n",
+    "        print(f\"\\n[Epoch {epoch:03d}] Train Loss={tr_loss:.3f} | Val Loss={va_loss:.3f}\")\n",
+    "        for k, v in metrics.items():\n",
+    "            print(\n",
+    "                f\"  {k:7s} | \"\n",
+    "                f\"Subtype AUC={v['Subtype_AUC']:.3f} | \"\n",
+    "                f\"TNM AUC={v['TNM_AUC_macro']:.3f} | \"\n",
+    "                f\"DFS C-index={v['DFS_C_index']:.3f} | \"\n",
+    "                f\"OS C-index={v['OS_C_index']:.3f}\"\n",
+    "            )\n",
+    "\n",
+    "        # early stopping\n",
+    "        if va_loss < best_val_loss:\n",
+    "            best_val_loss = va_loss\n",
+    "            wait = 0\n",
+    "            torch.save(model.state_dict(), os.path.join(SAVE_DIR, \"best_model.pt\"))\n",
+    "            print(\"  ✓ Best model updated\")\n",
+    "        else:\n",
+    "            wait += 1\n",
+    "            if wait >= PATIENCE:\n",
+    "                print(\"\\n⏹ Early stopping triggered\")\n",
+    "                break\n",
+    "\n",
+    "    # -------------------------\n",
+    "    # Inference (best model)\n",
+    "    # -------------------------\n",
+    "    print(\"\\nRunning inference with best model...\")\n",
+    "    model.load_state_dict(torch.load(os.path.join(SAVE_DIR, \"best_model.pt\"), map_location=DEVICE))\n",
+    "\n",
+    "    for split in [\"train\", \"val\", \"test\"]:\n",
+    "        out = run_epoch_verbose(model, loaders[split], optimizer, DEVICE, train=False)\n",
+    "        (\n",
+    "            loss,\n",
+    "            sy, ss,\n",
+    "            ty, ts,\n",
+    "            trt,\n",
+    "            dfs_r, dfs_t, dfs_e,\n",
+    "            os_r,  os_t,  os_e,\n",
+    "            dfs_log, os_log\n",
+    "        ) = out\n",
+    "\n",
+    "        # classification\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"subtype_{split}_labels.npy\"), sy)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"subtype_{split}_scores.npy\"), ss)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"tnm_{split}_labels.npy\"), ty)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"tnm_{split}_scores.npy\"), ts)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"treatment_{split}.npy\"), trt)\n",
+    "\n",
+    "        # survival (cox risk + time/event)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"dfs_{split}_risk.npy\"),  dfs_r)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"dfs_{split}_time.npy\"),  dfs_t)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"dfs_{split}_event.npy\"), dfs_e)\n",
+    "\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"os_{split}_risk.npy\"),   os_r)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"os_{split}_time.npy\"),   os_t)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"os_{split}_event.npy\"),  os_e)\n",
+    "\n",
+    "        # 1y/3y/5y logits (optional, for AUC at specific horizons)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"dfs_{split}_logits_1y3y5y.npy\"), dfs_log)\n",
+    "        np.save(os.path.join(SAVE_DIR, f\"os_{split}_logits_1y3y5y.npy\"),  os_log)\n",
+    "\n",
+    "        print(f\"{split:5s} | loss={loss:.3f} | Immune={np.sum(trt==0)} Chemo={np.sum(trt==1)}\")\n",
+    "\n",
+    "    print(\"\\n✓ Inference completed. Results saved.\")\n",
+    "\n",
+    "    # -------------------------\n",
+    "    # Figure: Immune/Chemo KM + HR\n",
+    "    # -------------------------\n",
+    "    print(\"\\nGenerating Figure  (KM + HR) ...\")\n",
+    "    generate_figure_from_saved(result_dir=SAVE_DIR, fig_dir=FIG_DIR, which_split=(\"val\", \"test\"))\n",
+    "    print(\"✓ Figure  done. Files saved under ./figures\")\n",
+    "\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    main()\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "VGG",
+   "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.6.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

test.py

@@ -0,0 +1,173 @@
+# test_mm_dls.py
+# =========================================================
+# 🔍 Minimal test for MM-DLS pipeline
+# - CUDA
+# - forward / loss
+# - pandas / lifelines (GLIBCXX check)
+# =========================================================
+
+import os
+import sys
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Subset
+
+import pandas as pd
+from lifelines import KaplanMeierFitter
+from lifelines.utils import concordance_index
+
+# ---------------------------------------------------------
+# Project path
+# ---------------------------------------------------------
+PROJECT_ROOT = os.path.abspath(".")
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+# ---------------------------------------------------------
+# Imports from mm_dls
+# ---------------------------------------------------------
+from mm_dls.HierMM_DLS import HierMM_DLS
+from mm_dls.CoxphLoss import CoxPHLoss
+from mm_dls.PatientDataset import PatientDataset
+
+
+# =========================================================
+# Basic config (VERY SMALL)
+# =========================================================
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print("Using device:", DEVICE)
+
+BATCH_SIZE = 2
+NUM_SUBTYPES = 2
+NUM_TNM = 3
+N_SLICES = 30
+IMG_SIZE = 224
+
+
+# =========================================================
+# Test Dataset Loader
+# =========================================================
+def get_test_loader():
+    dataset = PatientDataset(
+        data_root="/path/to/DATA_ROOT",
+        clinical_csv="/path/to/clinical.csv",
+        radiomics_npy="/path/to/radiomics.npy",
+        pet_npy="/path/to/pet.npy",
+        n_slices=N_SLICES,
+        img_size=IMG_SIZE,
+    )
+
+    # 🔑 只取前 8 个样本
+    idx = list(range(min(8, len(dataset))))
+    subset = Subset(dataset, idx)
+
+    loader = DataLoader(
+        subset,
+        batch_size=BATCH_SIZE,
+        shuffle=False,
+        num_workers=2,
+    )
+    return loader
+
+
+# =========================================================
+# One forward + loss
+# =========================================================
+def test_forward_and_loss():
+    print("\n[TEST] Forward + Loss")
+
+    loader = get_test_loader()
+    model = HierMM_DLS(NUM_SUBTYPES, NUM_TNM).to(DEVICE)
+
+    ce  = nn.CrossEntropyLoss()
+    bce = nn.BCEWithLogitsLoss()
+    cox = CoxPHLoss()
+
+    model.eval()
+
+    for batch in loader:
+        assert len(batch) == 19, f"Dataset must return 19 items, got {len(batch)}"
+
+        (
+            pid, lesion, space, rad, pet, cli,
+            y_sub, y_tnm,
+            dfs_t, dfs_e,
+            os_t, os_e,
+            dfs1, dfs3, dfs5,
+            os1, os3, os5,
+            treatment
+        ) = batch
+
+        lesion, space = lesion.to(DEVICE), space.to(DEVICE)
+        rad, pet, cli = rad.to(DEVICE), pet.to(DEVICE), cli.to(DEVICE)
+        y_sub, y_tnm  = y_sub.to(DEVICE), y_tnm.to(DEVICE)
+        dfs_t, dfs_e = dfs_t.to(DEVICE), dfs_e.to(DEVICE)
+        os_t,  os_e  = os_t.to(DEVICE),  os_e.to(DEVICE)
+
+        dfs_y = torch.stack([dfs1, dfs3, dfs5], dim=1).to(DEVICE)
+        os_y  = torch.stack([os1,  os3,  os5 ], dim=1).to(DEVICE)
+
+        with torch.no_grad():
+            sub_l, tnm_l, dfs_r, os_r, dfs_log, os_log = model(
+                lesion, space, rad, pet, cli
+            )
+
+            loss = (
+                ce(sub_l, y_sub) +
+                ce(tnm_l, y_tnm) +
+                cox(dfs_r, dfs_t, dfs_e) +
+                cox(os_r,  os_t,  os_e) +
+                bce(dfs_log, dfs_y) +
+                bce(os_log,  os_y)
+            )
+
+        print("  ✓ Forward OK | Loss =", float(loss))
+        break
+
+
+# =========================================================
+# Test pandas + lifelines (GLIBCXX killer)
+# =========================================================
+def test_pandas_lifelines():
+    print("\n[TEST] pandas + lifelines")
+
+    # fake survival data
+    time  = np.array([10, 12, 8, 20, 15, 25])
+    event = np.array([1, 1, 0, 1, 0, 0])
+    risk  = np.array([0.9, 0.8, 0.2, 1.2, 0.3, 0.4])
+
+    # pandas
+    df = pd.DataFrame({
+        "time": time,
+        "event": event,
+        "risk": risk
+    })
+
+    print("  pandas OK:", df.shape)
+
+    # C-index
+    cidx = concordance_index(df["time"], -df["risk"], df["event"])
+    print("  C-index =", round(cidx, 3))
+
+    # KM
+    kmf = KaplanMeierFitter()
+    kmf.fit(df["time"], event_observed=df["event"])
+    surv_10 = kmf.predict(10)
+
+    print("  KM survival@10 =", float(surv_10))
+    print("  ✓ lifelines OK")
+
+
+# =========================================================
+# Main
+# =========================================================
+if __name__ == "__main__":
+    print("\n==============================")
+    print(" MM-DLS TEST START ")
+    print("==============================")
+
+    test_forward_and_loss()
+    test_pandas_lifelines()
+
+    print("\n✅ ALL TESTS PASSED")

train_patient_model.py

@@ -0,0 +1,576 @@
+import sys
+sys.path.insert(0, "/export/home/daifang/lunghospital/MM-DLS-master/MM-DLS-master")
+# main.py
+import os
+import sys
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, random_split
+
+from sklearn.metrics import roc_auc_score, accuracy_score
+from sklearn.preprocessing import label_binarize
+
+import pandas as pd
+import matplotlib.pyplot as plt
+from lifelines import KaplanMeierFitter, CoxPHFitter
+from lifelines.statistics import multivariate_logrank_test
+from lifelines.utils import concordance_index
+from sklearn.metrics import brier_score_loss
+from scipy.stats import norm
+
+
+
+# =========================================================
+# Project path (IMPORTANT for Jupyter / HPC)
+# =========================================================
+PROJECT_ROOT = os.path.abspath(".")
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+# =========================================================
+# imports:  mm_dls/ 
+# =========================================================
+def _import_modules():
+
+    from mm_dls.HierMM_DLS import HierMM_DLS
+    from mm_dls.FakePatientDataset import FakePatientDataset
+    from mm_dls.CoxphLoss import CoxPHLoss
+    return HierMM_DLS, FakePatientDataset, CoxPHLoss
+
+
+HierMM_DLS, FakePatientDataset, CoxPHLoss = _import_modules()
+
+
+# =========================
+# Training configuration
+# =========================
+EPOCHS        = 300
+PATIENCE      = 8
+BATCH_SIZE    = 4
+LR            = 1e-4
+WEIGHT_DECAY  = 1e-5
+
+# =========================
+# Task definition
+# =========================
+NUM_SUBTYPES  = 2        # e.g., LUAD vs LUSC
+NUM_TNM       = 3        # Stage I–II / III / IV
+
+# =========================
+# Image settings
+# =========================
+N_SLICES      = 30       # max slices per patient
+IMG_SIZE      = 224
+
+
+SAVE_DIR = "./results"
+FIG_DIR  = "./figures"
+os.makedirs(SAVE_DIR, exist_ok=True)
+os.makedirs(FIG_DIR, exist_ok=True)
+
+# -------------------------
+# GPU (force cuda:1)
+# -------------------------
+assert torch.cuda.is_available(), "CUDA not available"
+DEVICE = torch.device("cuda:1")
+torch.cuda.set_device(DEVICE)
+print("Using device:", DEVICE)
+
+
+# =========================================================
+# Core utils
+# =========================================================
+def _sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+def _ensure_numpy(x):
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    return x
+
+def _risk_to_groups(risk, q=(1/3, 2/3), labels=("Low", "Mediate", "High")):
+    """
+    Convert continuous risk into 3 groups by tertiles.
+    """
+    r = np.asarray(risk).reshape(-1)
+    t1, t2 = np.quantile(r, q[0]), np.quantile(r, q[1])
+    out = np.full(len(r), labels[1], dtype=object)
+    out[r <= t1] = labels[0]
+    out[r >= t2] = labels[2]
+    return out
+
+def _evaluate_survival_metrics(time, event, risk, time_point=30):
+    """
+    C-index + Brier at a fixed time point.
+    risk: higher => earlier event, so use -risk in concordance_index.
+    """
+    time = np.asarray(time).reshape(-1)
+    event = np.asarray(event).reshape(-1).astype(int)
+    risk = np.asarray(risk).reshape(-1)
+
+    c_index = concordance_index(time, -risk, event)
+
+    # Brier: predict survival at time_point using a monotonic transform of risk (proxy)
+    # This is a "proxy" survival probability for demo/debug; replace with proper survival model if needed.
+    y_true = (time > time_point).astype(int)  # 1 means survived beyond time_point
+    # map risk into [0,1] survival prob proxy: higher risk => lower survival prob
+    y_prob = 1 - (risk - risk.min()) / (risk.max() - risk.min() + 1e-8)
+    brier = brier_score_loss(y_true, y_prob)
+
+    return float(c_index), float(brier)
+
+
+# =========================================================
+# One epoch (train / eval)
+# =========================================================
+def run_epoch_verbose(model, loader, optimizer, device, train=True):
+    ce  = nn.CrossEntropyLoss()
+    bce = nn.BCEWithLogitsLoss(reduction="none")
+    cox = CoxPHLoss()
+
+    model.train() if train else model.eval()
+
+    losses = []
+
+    # classification
+    sub_y_all, sub_s_all = [], []
+    tnm_y_all, tnm_s_all = [], []
+    treat_all = []
+
+    # survival (cox risk + time/event)
+    dfs_r_all, dfs_t_all, dfs_e_all = [], [], []
+    os_r_all,  os_t_all,  os_e_all  = [], [], []
+
+    # survival 1y/3y/5y logits (optional save)
+    dfs_log_all, os_log_all = [], []
+
+    for batch in loader:
+        # NOTE: dataset must return 19 items including treatment
+        if len(batch) != 19:
+            raise ValueError(f"Batch length mismatch: expected 19, got {len(batch)}. "
+                             f"Please ensure Dataset __getitem__ returns treatment as the 19th item.")
+
+        (
+            pid, lesion, space, rad, pet, cli,
+            y_sub, y_tnm,
+            dfs_t, dfs_e,
+            os_t, os_e,
+            dfs1, dfs3, dfs5,
+            os1, os3, os5,
+            treatment
+        ) = batch
+
+        lesion, space = lesion.to(device), space.to(device)
+        rad, pet, cli = rad.to(device), pet.to(device), cli.to(device)
+        y_sub, y_tnm  = y_sub.to(device), y_tnm.to(device)
+        dfs_t, dfs_e = dfs_t.to(device), dfs_e.to(device)
+        os_t,  os_e  = os_t.to(device),  os_e.to(device)
+        treatment = treatment.to(device)
+
+        dfs_y = torch.stack([dfs1, dfs3, dfs5], dim=1).to(device)
+        os_y  = torch.stack([os1,  os3,  os5 ], dim=1).to(device)
+
+        with torch.set_grad_enabled(train):
+            sub_l, tnm_l, dfs_r, os_r, dfs_log, os_log = model(
+                lesion, space, rad, pet, cli
+            )
+
+            loss = (
+                ce(sub_l, y_sub) +
+                ce(tnm_l, y_tnm) +
+                cox(dfs_r, dfs_t, dfs_e) +
+                cox(os_r,  os_t,  os_e) +
+                bce(dfs_log, dfs_y).mean() +
+                bce(os_log,  os_y ).mean()
+            )
+
+            if train:
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        losses.append(loss.item())
+
+        # ----- Collect predictions -----
+        sub_prob = torch.softmax(sub_l, dim=1)[:, 1]     # subtype prob
+        tnm_prob = torch.softmax(tnm_l, dim=1)           # [B,3]
+
+        sub_s_all.append(_ensure_numpy(sub_prob))
+        sub_y_all.append(_ensure_numpy(y_sub))
+
+        tnm_s_all.append(_ensure_numpy(tnm_prob))
+        tnm_y_all.append(_ensure_numpy(y_tnm))
+
+        treat_all.append(_ensure_numpy(treatment))
+
+        # survival
+        dfs_r_all.append(_ensure_numpy(dfs_r))
+        dfs_t_all.append(_ensure_numpy(dfs_t))
+        dfs_e_all.append(_ensure_numpy(dfs_e))
+
+        os_r_all.append(_ensure_numpy(os_r))
+        os_t_all.append(_ensure_numpy(os_t))
+        os_e_all.append(_ensure_numpy(os_e))
+
+        dfs_log_all.append(_ensure_numpy(dfs_log))
+        os_log_all.append(_ensure_numpy(os_log))
+
+    return (
+        float(np.mean(losses)),
+
+        np.concatenate(sub_y_all),
+        np.concatenate(sub_s_all),
+
+        np.concatenate(tnm_y_all),
+        np.concatenate(tnm_s_all),
+
+        np.concatenate(treat_all),
+
+        np.concatenate(dfs_r_all),
+        np.concatenate(dfs_t_all),
+        np.concatenate(dfs_e_all),
+
+        np.concatenate(os_r_all),
+        np.concatenate(os_t_all),
+        np.concatenate(os_e_all),
+
+        np.concatenate(dfs_log_all, axis=0),  # [N,3]
+        np.concatenate(os_log_all, axis=0),   # [N,3]
+    )
+
+
+# =========================================================
+# Evaluation by cohort (classification + survival)
+# =========================================================
+def evaluate_by_treatment(sub_y, sub_s, tnm_y, tnm_s, treat,
+                          dfs_r, dfs_t, dfs_e, os_r, os_t, os_e):
+    results = {}
+
+    cohorts = {
+        "All": np.ones_like(treat, dtype=bool),
+        "Immune": treat == 0,
+        "Chemo":  treat == 1,
+    }
+
+    for name, mask in cohorts.items():
+        if mask.sum() < 10:
+            continue
+
+        res = {}
+
+        # Subtype (binary)
+        res["Subtype_AUC"] = roc_auc_score(sub_y[mask], sub_s[mask])
+        res["Subtype_ACC"] = accuracy_score(sub_y[mask], (sub_s[mask] > 0.5).astype(int))
+
+        # TNM (multiclass macro AUC + ACC)
+        tnm_bin = label_binarize(tnm_y[mask], classes=[0, 1, 2])
+        res["TNM_AUC_macro"] = roc_auc_score(
+            tnm_bin, tnm_s[mask], average="macro", multi_class="ovr"
+        )
+        res["TNM_ACC"] = accuracy_score(
+            tnm_y[mask], np.argmax(tnm_s[mask], axis=1)
+        )
+
+        # Survival
+        dfs_c, dfs_b = _evaluate_survival_metrics(dfs_t[mask], dfs_e[mask], dfs_r[mask], time_point=30)
+        os_c,  os_b  = _evaluate_survival_metrics(os_t[mask],  os_e[mask],  os_r[mask],  time_point=30)
+
+        res["DFS_C_index"] = dfs_c
+        res["DFS_Brier_30m"] = dfs_b
+        res["OS_C_index"] = os_c
+        res["OS_Brier_30m"] = os_b
+
+        results[name] = res
+
+    return results
+
+
+# =========================================================
+# Figure 7: KM + HR (per cohort, per endpoint)
+# =========================================================
+def plot_km_curve_with_hr(df, title, save_prefix):
+    """
+    df must contain columns: time, event, group (Low/Mediate/High)
+    """
+    kmf = KaplanMeierFitter()
+    fig, ax = plt.subplots(figsize=(8, 6), facecolor="white")
+    ax.set_facecolor("white")
+
+    colors = {"Low": "#91c7ae", "Mediate": "#f7b977", "High": "#d87c7c"}
+    groups = ["Low", "Mediate", "High"]
+
+    # plot KM
+    lines = {}
+    at_risk_table = []
+    times = np.arange(0, 70, 10)
+
+    for g in groups:
+        m = df["group"] == g
+        if m.sum() == 0:
+            continue
+
+        kmf.fit(df.loc[m, "time"], event_observed=df.loc[m, "event"], label=g)
+        kmf.plot_survival_function(
+            ax=ax, ci_show=True, linewidth=2, color=colors[g], marker="+"
+        )
+        lines[g] = ax.get_lines()[-1]
+
+        at_risk_table.append([np.sum(df.loc[m, "time"] >= t) for t in times])
+
+    # legend
+    handles = [lines[g] for g in groups if g in lines]
+    labels = ["Low", "Medium", "High"][:len(handles)]
+    ax.legend(handles, labels, title="Groups", loc="upper right",
+              frameon=True, framealpha=0.5, fontsize=12, title_fontsize=12)
+
+    # at risk numbers (optional, matches your style)
+    if len(at_risk_table) == 3:
+        low, mid, high = at_risk_table
+        for i, t in enumerate(times):
+            ax.text(t, -0.38, str(low[i]),  color="#207f4c", fontsize=14, ha="center")
+            ax.text(t, -0.48, str(mid[i]),  color="#fca106", fontsize=14, ha="center")
+            ax.text(t, -0.58, str(high[i]), color="#cc163a", fontsize=14, ha="center")
+
+        ax.text(-1,  -0.28, "Number at risk", color="black", ha="center", fontsize=14)
+        ax.text(-10, -0.38, "Low",    color="#207f4c", fontsize=14)
+        ax.text(-10, -0.48, "Medium", color="#fca106", fontsize=14)
+        ax.text(-10, -0.58, "High",   color="#cc163a", fontsize=14)
+
+    # Cox HR + p-values
+    df2 = df.copy()
+    df2["group_code"] = df2["group"].map({"Low": 0, "Mediate": 1, "High": 2})
+    cph = CoxPHFitter()
+    cph.fit(df2[["time", "event", "group_code"]], duration_col="time", event_col="event")
+
+    coef = float(cph.params_["group_code"])
+    se   = float(cph.standard_errors_["group_code"])
+
+    hr_med_vs_low  = np.exp(coef * 1)
+    hr_high_vs_low = np.exp(coef * 2)
+
+    z_med  = (coef * 1) / se
+    p_med  = 2 * (1 - norm.cdf(abs(z_med)))
+
+    z_high = (coef * 2) / se
+    p_high = 2 * (1 - norm.cdf(abs(z_high)))
+
+    # logrank
+    res_lr = multivariate_logrank_test(df2["time"], df2["group"], df2["event"])
+
+    # C-index + brier (proxy)
+    c_index, brier = _evaluate_survival_metrics(df2["time"].values, df2["event"].values,
+                                                df2["group_code"].values, time_point=30)
+
+    ax.text(25, 0.46, f"P(log-rank)={res_lr.p_value:.3f}", fontsize=12)
+    ax.text(25, 0.36, f"C-index={c_index:.3f}", fontsize=12)
+    ax.text(25, 0.26, f"Brier(30m)={brier:.3f}", fontsize=12)
+    ax.text(25, 0.16, f"HR Intermediate vs Low = {hr_med_vs_low:.2f}, P={p_med:.3f}", fontsize=12)
+    ax.text(25, 0.06, f"HR High vs Low = {hr_high_vs_low:.2f}, P={p_high:.3f}", fontsize=12)
+
+    # cosmetics
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.set_title(title, fontsize=14)
+    ax.set_xlabel("Time since treatment start (months)", fontsize=14)
+    ax.set_ylabel("Survival probability", fontsize=14)
+    ax.set_ylim(0, 1.05)
+    ax.grid(alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(save_prefix + ".png", dpi=600, bbox_inches="tight")
+    plt.savefig(save_prefix + ".pdf", dpi=600, bbox_inches="tight")
+    plt.close()
+    return save_prefix
+
+
+def generate_figure_from_saved(result_dir=SAVE_DIR, fig_dir=FIG_DIR, which_split=("val", "test")):
+    """
+    Load saved dfs/os arrays and generate KM+HR for Immune/Chemo separately.
+    """
+    os.makedirs(fig_dir, exist_ok=True)
+
+    for split in which_split:
+        # load arrays
+        trt = np.load(os.path.join(result_dir, f"treatment_{split}.npy"))
+
+        dfs_r = np.load(os.path.join(result_dir, f"dfs_{split}_risk.npy"))
+        dfs_t = np.load(os.path.join(result_dir, f"dfs_{split}_time.npy"))
+        dfs_e = np.load(os.path.join(result_dir, f"dfs_{split}_event.npy"))
+
+        os_r  = np.load(os.path.join(result_dir, f"os_{split}_risk.npy"))
+        os_t  = np.load(os.path.join(result_dir, f"os_{split}_time.npy"))
+        os_e  = np.load(os.path.join(result_dir, f"os_{split}_event.npy"))
+
+        for cohort_name, mask in {
+            "Immune": trt == 0,
+            "Chemo":  trt == 1
+        }.items():
+            if mask.sum() < 20:
+                print(f"[Figure7] Skip {split}-{cohort_name}: too few samples ({mask.sum()})")
+                continue
+
+            # DFS groups
+            dfs_group = _risk_to_groups(dfs_r[mask])
+            df_dfs = pd.DataFrame({
+                "time": dfs_t[mask],
+                "event": dfs_e[mask].astype(int),
+                "group": dfs_group
+            })
+
+            # OS groups
+            os_group = _risk_to_groups(os_r[mask])
+            df_os = pd.DataFrame({
+                "time": os_t[mask],
+                "event": os_e[mask].astype(int),
+                "group": os_group
+            })
+
+            # save CSV (optional, for reproducibility)
+            df_dfs.to_csv(os.path.join(result_dir, f"dfs_{split}_{cohort_name}.csv"), index=False)
+            df_os.to_csv(os.path.join(result_dir, f"os_{split}_{cohort_name}.csv"), index=False)
+
+            # plot
+            plot_km_curve_with_hr(
+                df_dfs,
+                title=f"Disease-Free Survival (DFS) — Kaplan-Meier Curves\n{cohort_name} {split} set (n={mask.sum()})",
+                save_prefix=os.path.join(fig_dir, f"Figure7_DFS_{cohort_name}_{split}")
+            )
+            plot_km_curve_with_hr(
+                df_os,
+                title=f"Overall Survival (OS) — Kaplan-Meier Curves\n{cohort_name} {split} set (n={mask.sum()})",
+                save_prefix=os.path.join(fig_dir, f"Figure7_OS_{cohort_name}_{split}")
+            )
+
+    print("✔ Figure 7 generated (DFS/OS KM + HR) for Immune/Chemo.")
+
+
+# =========================================================
+# Main
+# =========================================================
+def main():
+    # -------------------------
+    # Dataset (must return treatment as 19th item)
+    # -------------------------
+    from mm_dls.PatientDataset import PatientDataset
+
+    dataset = PatientDataset(
+        data_root="/path/to/DATA_ROOT",
+        clinical_csv="/path/to/clinical.csv",
+        radiomics_npy="/path/to/radiomics.npy",
+        pet_npy="/path/to/pet.npy",
+        n_slices=N_SLICES,
+        img_size=IMG_SIZE,
+    )
+
+
+    n_train = int(0.6 * len(dataset))
+    n_val   = int(0.2 * len(dataset))
+    n_test  = len(dataset) - n_train - n_val
+
+    train_set, val_set, test_set = random_split(dataset, [n_train, n_val, n_test])
+
+    loaders = {
+        "train": DataLoader(train_set, BATCH_SIZE, shuffle=True,  num_workers=4),
+        "val":   DataLoader(val_set,   BATCH_SIZE, shuffle=False, num_workers=4),
+        "test":  DataLoader(test_set,  BATCH_SIZE, shuffle=False, num_workers=4),
+    }
+
+    # -------------------------
+    # Model
+    # -------------------------
+    model = HierMM_DLS(NUM_SUBTYPES, NUM_TNM).to(DEVICE)
+    optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
+
+    best_val_loss = 1e9
+    wait = 0
+
+    # -------------------------
+    # Training
+    # -------------------------
+    for epoch in range(1, EPOCHS + 1):
+        tr = run_epoch_verbose(model, loaders["train"], optimizer, DEVICE, train=True)
+        va = run_epoch_verbose(model, loaders["val"],   optimizer, DEVICE, train=False)
+
+        tr_loss = tr[0]
+        va_loss = va[0]
+
+        # unpack val for metrics
+        _, sy, ss, ty, ts, trt, dfs_r, dfs_t, dfs_e, os_r, os_t, os_e, _, _ = va
+        metrics = evaluate_by_treatment(sy, ss, ty, ts, trt, dfs_r, dfs_t, dfs_e, os_r, os_t, os_e)
+
+        print(f"\n[Epoch {epoch:03d}] Train Loss={tr_loss:.3f} | Val Loss={va_loss:.3f}")
+        for k, v in metrics.items():
+            print(
+                f"  {k:7s} | "
+                f"Subtype AUC={v['Subtype_AUC']:.3f} | "
+                f"TNM AUC={v['TNM_AUC_macro']:.3f} | "
+                f"DFS C-index={v['DFS_C_index']:.3f} | "
+                f"OS C-index={v['OS_C_index']:.3f}"
+            )
+
+        # early stopping
+        if va_loss < best_val_loss:
+            best_val_loss = va_loss
+            wait = 0
+            torch.save(model.state_dict(), os.path.join(SAVE_DIR, "best_model.pt"))
+            print("  ✓ Best model updated")
+        else:
+            wait += 1
+            if wait >= PATIENCE:
+                print("\n⏹ Early stopping triggered")
+                break
+
+    # -------------------------
+    # Inference (best model)
+    # -------------------------
+    print("\nRunning inference with best model...")
+    model.load_state_dict(torch.load(os.path.join(SAVE_DIR, "best_model.pt"), map_location=DEVICE))
+
+    for split in ["train", "val", "test"]:
+        out = run_epoch_verbose(model, loaders[split], optimizer, DEVICE, train=False)
+        (
+            loss,
+            sy, ss,
+            ty, ts,
+            trt,
+            dfs_r, dfs_t, dfs_e,
+            os_r,  os_t,  os_e,
+            dfs_log, os_log
+        ) = out
+
+        # classification
+        np.save(os.path.join(SAVE_DIR, f"subtype_{split}_labels.npy"), sy)
+        np.save(os.path.join(SAVE_DIR, f"subtype_{split}_scores.npy"), ss)
+        np.save(os.path.join(SAVE_DIR, f"tnm_{split}_labels.npy"), ty)
+        np.save(os.path.join(SAVE_DIR, f"tnm_{split}_scores.npy"), ts)
+        np.save(os.path.join(SAVE_DIR, f"treatment_{split}.npy"), trt)
+
+        # survival (cox risk + time/event)
+        np.save(os.path.join(SAVE_DIR, f"dfs_{split}_risk.npy"),  dfs_r)
+        np.save(os.path.join(SAVE_DIR, f"dfs_{split}_time.npy"),  dfs_t)
+        np.save(os.path.join(SAVE_DIR, f"dfs_{split}_event.npy"), dfs_e)
+
+        np.save(os.path.join(SAVE_DIR, f"os_{split}_risk.npy"),   os_r)
+        np.save(os.path.join(SAVE_DIR, f"os_{split}_time.npy"),   os_t)
+        np.save(os.path.join(SAVE_DIR, f"os_{split}_event.npy"),  os_e)
+
+        # 1y/3y/5y logits (optional, for AUC at specific horizons)
+        np.save(os.path.join(SAVE_DIR, f"dfs_{split}_logits_1y3y5y.npy"), dfs_log)
+        np.save(os.path.join(SAVE_DIR, f"os_{split}_logits_1y3y5y.npy"),  os_log)
+
+        print(f"{split:5s} | loss={loss:.3f} | Immune={np.sum(trt==0)} Chemo={np.sum(trt==1)}")
+
+    print("\n✓ Inference completed. Results saved.")
+
+    # -------------------------
+    # Figure: Immune/Chemo KM + HR
+    # -------------------------
+    print("\nGenerating Figure  (KM + HR) ...")
+    generate_figure_from_saved(result_dir=SAVE_DIR, fig_dir=FIG_DIR, which_split=("val", "test"))
+    print("✓ Figure  done. Files saved under ./figures")
+
+
+if __name__ == "__main__":
+    main()
+