Add CheXlocalize_Dataset for CheXpert official val/test + segmentation masks

README.md

@@ -190,6 +190,10 @@ d_vin = xrv.datasets.VinBrain_Dataset(imgpath=".../train",
 
 # National Library of Medicine Tuberculosis Datasets. https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4256233/
 d_nlmtb = xrv.datasets.NLMTB_Dataset(imgpath="path to MontgomerySet or ChinaSet_AllFiles")
+
+# CheXlocalize: official CheXpert val/test images, blinded test labels, and radiologist segmentations. https://doi.org/10.1038/s42256-022-00536-x
+d_chexlocalize = xrv.datasets.CheXlocalize_Dataset(imgpath="path to CheXpert val or test image folder",
+                                                   csvpath="path to val_labels.csv or test_labels.csv")
 ```
 
 ## Dataset fields
@@ -263,6 +267,7 @@ Masks are available in the following datasets:
 xrv.datasets.RSNA_Pneumonia_Dataset() # for Lung Opacity
 xrv.datasets.SIIM_Pneumothorax_Dataset() # for Pneumothorax
 xrv.datasets.NIH_Dataset() # for Cardiomegaly, Mass, Effusion, ...
+xrv.datasets.CheXlocalize_Dataset() # for 10 pathologies, from radiologist ground-truth segmentations
 ```
 
 Example usage:

scripts/xray_chexlocalize.ipynb

@@ -0,0 +1,117 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os,sys\n",
+    "sys.path.insert(0,\"..\")\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torchxrayvision as xrv"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "`CheXlocalize_Dataset` follows the same `pathology_masks` interface as `NIH_Dataset`, `VinBrain_Dataset`, and `SIIM_Pneumothorax_Dataset`: `sample[\"pathology_masks\"]` is a sparse dict keyed by pathology index, each value a `(1, H, W)` array."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_path = \"/home/groups/akshaysc/joecohen/CheXpert\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_sample_with_masks(sample, df):\n",
+    "    width = len(sample[\"pathology_masks\"])\n",
+    "    fig, axs = plt.subplots(1, max(2,1+width), sharey=True, figsize=(3+3*width,3))\n",
+    "    axs[0].imshow(sample[\"img\"][0], cmap=\"Greys_r\");\n",
+    "    axs[0].set_title(\"idx:\" + str(sample[\"idx\"]))\n",
+    "    for i, patho in enumerate(sample[\"pathology_masks\"].keys()):\n",
+    "        axs[i+1].imshow(sample[\"img\"][0], cmap=\"Greys_r\");\n",
+    "        axs[i+1].imshow(sample[\"pathology_masks\"][patho][0]+1, alpha=0.5);\n",
+    "        axs[i+1].set_title(df.pathologies[patho])\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_chexlocalize = xrv.datasets.CheXlocalize_Dataset(\n",
+    "    imgpath=os.path.join(dataset_path, \"CheXpert-v1.0/test\"),\n",
+    "    csvpath=os.path.join(dataset_path, \"test_labels.csv\"),\n",
+    "    pathology_masks=True,\n",
+    "    segmentation_jsonpath=os.path.join(dataset_path, \"gt_segmentations_test.json\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_chexlocalize.csv.has_masks.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for idx in np.where(d_chexlocalize.csv.has_masks)[0][:5]:\n",
+    "    sample = d_chexlocalize[idx]\n",
+    "    if len(sample[\"pathology_masks\"]) > 0:\n",
+    "        plot_sample_with_masks(sample, d_chexlocalize)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

tests/test_dataloaders.py

@@ -1,3 +1,4 @@
+import json
 import os
 import shutil
 import sys
@@ -323,3 +324,102 @@ def test_nih_dataset_patient_sex_column(tmp_path):
 
     assert d.csv["sex_male"].iloc[0]
     assert not d.csv["sex_female"].iloc[0]
+
+
+def _make_chexlocalize_test_csv(tmp_path, split="test"):
+    """A blinded-official-test-style CSV: no Sex/Age/Frontal-Lateral/AP-PA
+    columns, and a 'valid/' path prefix on val rows to test normalization."""
+    import pandas as pd
+
+    prefix = "valid/" if split == "val" else f"{split}/"
+    csv = pd.DataFrame({
+        "Path": [f"CheXpert-v1.0/{prefix}patient64622/study1/view1_frontal.jpg"],
+        "Atelectasis": [1],
+        "Cardiomegaly": [0],
+        "Consolidation": [np.nan],
+        "Edema": [-1],
+        "Pleural Effusion": [1],
+    })
+    csv_path = tmp_path / f"{split}_labels.csv"
+    csv.to_csv(csv_path, index=False)
+
+    img_dir = tmp_path / ("val" if split == "val" else split) / "patient64622" / "study1"
+    img_dir.mkdir(parents=True)
+    shutil.copyfile(test_jpg_img_file, img_dir / "view1_frontal.jpg")
+
+    return csv_path
+
+
+def test_chexlocalize_dataset_blinded_test_csv(tmp_path):
+    """CheX_Dataset raises NotImplementedError on a CSV path without 'train'
+    or 'valid' in it (as with CheXlocalize's test_labels.csv), and assumes
+    Sex/Age/Frontal-Lateral/AP-PA columns exist. CheXlocalize_Dataset must
+    load this blinded test CSV without either problem."""
+    csv_path = _make_chexlocalize_test_csv(tmp_path, split="test")
+
+    d = xrv.datasets.CheXlocalize_Dataset(imgpath=str(tmp_path), csvpath=str(csv_path), pathology_masks=False)
+
+    # AP/PA is genuinely unknown for the blinded test CSV — must not be
+    # guessed as "AP", and the frontal image must still survive the default
+    # views filter rather than being silently dropped as an unknown view.
+    assert d.csv["AP/PA"].iloc[0] == "Unknown"
+    assert len(d) == 1
+    sample = d[0]
+    assert "img" in sample
+    assert "lab" in sample
+    assert d.csv["patientid"].iloc[0] == "64622"
+
+    atelectasis_idx = d.pathologies.index("Atelectasis")
+    edema_idx = d.pathologies.index("Edema")
+    assert sample["lab"][atelectasis_idx] == 1
+    assert np.isnan(sample["lab"][edema_idx])  # uncertain (-1) becomes NaN
+
+
+def test_chexlocalize_dataset_valid_prefix_normalized_to_val(tmp_path):
+    """CheXlocalize's val CSV uses a 'valid/' path prefix, but images ship
+    under a 'val/' directory on disk — CheXlocalize_Dataset must normalize
+    this rather than fail to find the image."""
+    csv_path = _make_chexlocalize_test_csv(tmp_path, split="val")
+
+    d = xrv.datasets.CheXlocalize_Dataset(imgpath=str(tmp_path), csvpath=str(csv_path), pathology_masks=False)
+
+    assert len(d) == 1
+    sample = d[0]
+    assert "img" in sample
+
+
+def test_chexlocalize_dataset_segmentation_masks(tmp_path):
+    pycocotools = pytest.importorskip("pycocotools")
+    from pycocotools import mask as coco_mask
+
+    csv_path = _make_chexlocalize_test_csv(tmp_path, split="test")
+
+    raw_mask = np.zeros((256, 256), dtype=np.uint8, order="F")
+    raw_mask[64:128, 64:128] = 1
+    rle = coco_mask.encode(raw_mask)
+    rle["counts"] = rle["counts"].decode("ascii")
+
+    segmentations = {
+        "patient64622_study1_view1_frontal": {
+            "Atelectasis": rle,
+        }
+    }
+    seg_path = tmp_path / "gt_segmentations_test.json"
+    with open(seg_path, "w") as f:
+        json.dump(segmentations, f)
+
+    d = xrv.datasets.CheXlocalize_Dataset(
+        imgpath=str(tmp_path), csvpath=str(csv_path),
+        pathology_masks=True, segmentation_jsonpath=str(seg_path))
+
+    sample = d[0]
+    assert "pathology_masks" in sample
+
+    atelectasis_idx = d.pathologies.index("Atelectasis")
+    cardiomegaly_idx = d.pathologies.index("Cardiomegaly")
+    assert atelectasis_idx in sample["pathology_masks"]
+    assert sample["pathology_masks"][atelectasis_idx].sum() > 0
+    # Pathologies absent from the segmentation JSON get no entry (sparse dict,
+    # matching NIH_Dataset/VinBrain_Dataset/ObjectCXR_Dataset)
+    assert cardiomegaly_idx not in sample["pathology_masks"]
+    assert bool(d.csv["has_masks"].iloc[0])