changed docker files

Author: vaibhavbahel

src/baseline/v2/inference/docker/Dockerfile

@@ -2,18 +2,30 @@ FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
 
 WORKDIR /app
 
-# system deps
+# -----------------------------
+# System dependencies
+# -----------------------------
 RUN apt-get update && apt-get install -y \
     git \
     wget \
     && rm -rf /var/lib/apt/lists/*
 
-# python deps
+# -----------------------------
+# Python dependencies
+# -----------------------------
 COPY requirements.txt .
+
+# remove torchvision (not needed + removes warning)
+RUN pip uninstall -y torchvision || true
+
 RUN pip install --no-cache-dir -r requirements.txt
 
-# copy code
+# -----------------------------
+# Copy code
+# -----------------------------
 COPY . .
 
-# run inference
+# -----------------------------
+# Run inference
+# -----------------------------
 ENTRYPOINT ["python", "inference.py"]

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

@@ -14,7 +14,7 @@
 )
 
 from models.unet import build_model
-
+import time 
 
 # -----------------------------
 # ARGUMENTS
@@ -25,6 +25,9 @@
 parser.add_argument("--output", required=True, help="Path to save pseudo CT")
 args = parser.parse_args()
 
+total_start = time.time()
+
+
 INPUT_PATH = Path(args.input)
 OUTPUT_PATH = Path(args.output)
 
@@ -99,4 +102,7 @@
     str(OUTPUT_PATH)
 )
 
-print("Saved:", OUTPUT_PATH)
+print("Saved:", OUTPUT_PATH)
+
+total_end = time.time()
+print(f"Total runtime: {total_end - total_start:.2f} seconds")

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

@@ -1,3 +1,3 @@
-monai
+monai==1.3.0
 nibabel
 tqdm