added docker

Author: vaibhavbahel

src/baseline/v2/inference/docker/Dockerfile

@@ -0,0 +1,19 @@
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
+
+WORKDIR /app
+
+# system deps
+RUN apt-get update && apt-get install -y \
+    git \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+# python deps
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# copy code
+COPY . .
+
+# run inference
+ENTRYPOINT ["python", "inference.py"]

src/baseline/v2/inference/docker/inference.py

@@ -0,0 +1,102 @@
+import argparse
+import torch
+import nibabel as nib
+from pathlib import Path
+
+from monai.inferers import sliding_window_inference
+from monai.transforms import (
+    Compose,
+    LoadImaged,
+    EnsureChannelFirstd,
+    NormalizeIntensityd,
+    ConcatItemsd,
+    EnsureTyped,
+)
+
+from models.unet import build_model
+
+
+# -----------------------------
+# ARGUMENTS
+# -----------------------------
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--input", required=True, help="Path to nacpet.nii.gz")
+parser.add_argument("--output", required=True, help="Path to save pseudo CT")
+args = parser.parse_args()
+
+INPUT_PATH = Path(args.input)
+OUTPUT_PATH = Path(args.output)
+
+OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
+
+
+# -----------------------------
+# CONFIG
+# -----------------------------
+
+MODEL_PATH = "weights/best_model.pth"
+PATCH_SIZE = (192,192,192)
+SW_BATCH = 2
+OVERLAP = 0.75
+
+
+# -----------------------------
+# TRANSFORMS (PET only)
+# -----------------------------
+
+transforms = Compose([
+    LoadImaged(keys=["pet"]),
+    EnsureChannelFirstd(keys=["pet"]),
+    NormalizeIntensityd(keys=["pet"], nonzero=True, channel_wise=True),
+    ConcatItemsd(keys=["pet"], name="input"),
+    EnsureTyped(keys=["input"])
+])
+
+
+# -----------------------------
+# MODEL
+# -----------------------------
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
+
+model = build_model().to(device)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+model.eval()
+
+
+# -----------------------------
+# INFERENCE
+# -----------------------------
+
+data = {"pet": INPUT_PATH}
+
+data = transforms(data)
+
+x = data["input"].unsqueeze(0).to(device)
+
+with torch.no_grad():
+
+    pred = sliding_window_inference(
+        x,
+        PATCH_SIZE,
+        SW_BATCH,
+        model,
+        overlap=OVERLAP,
+        mode="gaussian",
+    )
+
+pred = pred.cpu().numpy()[0,0]
+
+# convert back to HU
+pred = pred * 3000 - 1000
+
+ref = nib.load(str(INPUT_PATH))
+
+nib.save(
+    nib.Nifti1Image(pred, ref.affine, ref.header),
+    str(OUTPUT_PATH)
+)
+
+print("Saved:", OUTPUT_PATH)

src/baseline/v2/inference/docker/requirments.txt

@@ -0,0 +1,3 @@
+monai
+nibabel
+tqdm