Files
patient_leakage_detection/scripts/visualizations/figure1.py
T
rpotter6298 35cbd9ac3c 2026001
2026-07-01 17:35:58 +02:00

72 lines
2.9 KiB
Python

#!/usr/bin/env python3
"""figure1.py — Sample CT images from the IQ-OTH/NCCD dataset, one per patient.
Usage: python scripts/visualizations/figure1.py [--tag TAG]"""
import os, sys, csv, argparse
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
os.path.abspath(__file__)))))
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
os.path.abspath(__file__))))
DATASET = os.path.join(os.path.dirname(ROOT), "The IQ-OTHNCCD lung cancer dataset")
MANIFEST = os.path.join(ROOT, "results", "simple_patient_manifest.csv")
PLOTS_DIR = os.path.join(ROOT, "plots")
ap = argparse.ArgumentParser()
ap.add_argument("--tag", default="", help="Append tag to filename")
ap.add_argument("--manifest", default=MANIFEST, help="Patient manifest CSV")
args = ap.parse_args()
tag = f"_{args.tag}" if args.tag else ""
# Load manifest to get per-patient images
patients = {"Benign": [], "Malignant": [], "Normal": []}
with open(args.manifest, newline="") as f:
reader = csv.DictReader(f)
img_col = "confirmed_images" if "confirmed_images" in reader.fieldnames else "images"
for row in reader:
cls = row.get("class", "")
if cls in patients:
imgs = row[img_col].split(";")
if imgs:
patients[cls].append((row["patient_id"], imgs[0])) # first image per patient
CLASS_DIR = {"Benign": "Bengin cases", "Malignant": "Malignant cases",
"Normal": "Normal cases"}
N_EXAMPLES = 3
fig, axes = plt.subplots(3, N_EXAMPLES, figsize=(8, 9))
for row, (cls_label, cls_dir) in enumerate(CLASS_DIR.items()):
# Pick first N_EXAMPLES patients for this class
selected = patients[cls_label][:N_EXAMPLES]
for col, (pid, short_name) in enumerate(selected):
ax = axes[row, col]
# Convert short name back to original filename
prefix = short_name[0]
num = int(short_name.split("_")[1])
cls_map = {"B": ("Bengin cases", "Bengin"), "M": ("Malignant cases", "Malignant"),
"N": ("Normal cases", "Normal")}
dir_name, file_prefix = cls_map[prefix]
fname = f"{file_prefix} case ({num}).jpg"
img_path = os.path.join(DATASET, dir_name, fname)
try:
img = Image.open(img_path).convert("L")
ax.imshow(img, cmap="gray")
except Exception as e:
ax.text(0.5, 0.5, f"error: {e}", ha="center", va="center", fontsize=7)
ax.set_xticks([]); ax.set_yticks([])
if col == 0:
ax.set_ylabel(cls_label, fontsize=10, rotation=0,
labelpad=20, va="center")
fig.suptitle("Figure 1 — Sample images from the IQ-OTH/NCCD dataset",
fontsize=12, y=1.02)
plt.tight_layout(rect=[0, 0, 1, 0.97])
out = os.path.join(PLOTS_DIR, "figure1", f"figure1{tag}.png")
os.makedirs(os.path.dirname(out), exist_ok=True)
plt.savefig(out, dpi=150)
plt.close()
print(f"Saved → {out}")