cxrmate-ed / dataset.py
anicolson's picture
Upload model
6f7f115 verified
raw
history blame
14.3 kB
import os
import struct
import lmdb
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from torchvision.io import decode_image, read_image
from data.mimic_cxr.dcm_processing import load_and_preprocess_dcm_uint16
from tools.mimic_iv.ed_cxr.records import EDCXRSubjectRecords
from tools.utils import mimic_cxr_image_path
# Ordered by oblique, lateral, AP, and then PA views so that PA views are closest in position to the generated tokens (and oblique is furtherest).
VIEW_ORDER = ['LPO', 'RAO', 'LAO', 'SWIMMERS', 'XTABLE LATERAL', 'LL', 'LATERAL', 'AP AXIAL', 'AP RLD', 'AP LLD', 'AP', 'PA RLD', 'PA LLD', 'PA']
class StudyIDEDStayIDSubset(Dataset):
"""
Study ID & ED stay ID subset. Examples are indexed by the study identifier.
Information from the ED module is added by finding the study_id that is within
the timespan of the stay_id for the subject_id. The history and indication
sections are also included.
"""
def __init__(
self,
mimic_iv_duckdb_path,
split,
dataset_dir=None,
max_images_per_study=None,
transforms=None,
images=True,
columns='study_id, dicom_id, subject_id, findings, impression',
and_condition='',
records=None,
study_id_inclusion_list=None,
return_images=True,
ed_module=True,
extension='jpg',
images_rocksdb_path=None,
jpg_lmdb_path=None,
jpg_rocksdb_path=None,
):
"""
Argument/s:
mimic_iv_duckdb_path - Path to MIMIC-IV DuckDB database.
split - 'train', 'validate', or 'test'.
dataset_dir - Dataset directory.
max_images_per_study - the maximum number of images per study.
transforms - torchvision transformations.
colour_space - PIL target colour space.
images - flag to return processed images.
columns - which columns to query on.
and_condition - AND condition to add to the SQL query.
records - MIMIC-IV records class instance.
study_id_inclusion_list - studies not in this list are excluded.
return_images - return CXR images for the study as tensors.
ed_module - use the ED module.
extension - 'jpg' or 'dcm'.
images_rocksdb_path - path to image RocksDB database.
jpg_lmdb_path - path to LMDB .jpg database.
jpg_rocksdb_path - path to RocksDB .jpg database.
"""
super(StudyIDEDStayIDSubset, self).__init__()
self.split = split
self.dataset_dir = dataset_dir
self.max_images_per_study = max_images_per_study
self.transforms = transforms
self.images = images
self.columns = columns
self.and_condition = and_condition
self.return_images = return_images
self.ed_module = ed_module
self.extension = extension
self.images_rocksdb_path = images_rocksdb_path
self.jpg_lmdb_path = jpg_lmdb_path
self.jpg_rocksdb_path = jpg_rocksdb_path
# If max images per study is not set:
self.max_images_per_study = float('inf') if self.max_images_per_study is None else self.max_images_per_study
assert self.extension == 'jpg' or self.extension == 'dcm'
if self.dataset_dir is not None and self.images_rocksdb_path is None:
if self.extension == 'jpg':
if 'physionet.org/files/mimic-cxr-jpg/2.0.0/files' not in self.dataset_dir:
self.dataset_dir = os.path.join(self.dataset_dir, 'physionet.org/files/mimic-cxr-jpg/2.0.0/files')
elif self.extension == 'dcm':
if 'physionet.org/files/mimic-cxr/2.0.0/files' not in self.dataset_dir:
self.dataset_dir = os.path.join(self.dataset_dir, 'physionet.org/files/mimic-cxr/2.0.0/files')
# Open the RocksDB images database:
if self.images_rocksdb_path is not None:
import rocksdb
# Define the column families:
column_families = {
b'shape': rocksdb.ColumnFamilyOptions(),
b'image': rocksdb.ColumnFamilyOptions(),
}
opts = rocksdb.Options()
opts.max_open_files = 1e+5
self.images_db = rocksdb.DB(self.images_rocksdb_path, opts, column_families=column_families, read_only=True)
self.shape_handle = self.images_db.get_column_family(b'shape')
self.image_handle = self.images_db.get_column_family(b'image')
self.shape_dtype = np.int32
self.image_dtype = np.uint16
# Prepare the RocksDB .jpg database:
if self.jpg_rocksdb_path is not None:
import rocksdb
opts = rocksdb.Options()
opts.max_open_files = 1e+5
self.images_db = rocksdb.DB(self.jpg_rocksdb_path, opts, read_only=True)
# Prepare the LMDB .jpg database:
if self.jpg_lmdb_path is not None:
print('Loading images using LMDB.')
# Map size:
map_size = int(0.65 * (1024 ** 4))
assert isinstance(map_size, int)
self.env = lmdb.open(self.jpg_lmdb_path, map_size=map_size, lock=False, readonly=True)
self.txn = self.env.begin(write=False)
self.records = EDCXRSubjectRecords(database_path=mimic_iv_duckdb_path) if records is None else records
query = f"""
SELECT {columns}
FROM mimic_cxr
WHERE split = '{split}'
{and_condition}
ORDER BY study_id
"""
# For multi-image, the study identifiers make up the training examples:
df = self.records.connect.sql(query).df()
# Drop studies that don't have a findings or impression section:
df = df.dropna(subset=['findings', 'impression'], how='any')
# This study has two rows in edstays (removed as it causes issues):
if self.ed_module:
df = df[df['study_id'] != 59128861]
# Exclude studies not in list:
if study_id_inclusion_list is not None:
df = df[df['study_id'].isin(study_id_inclusion_list)]
# Example study identifiers for the subset:
self.examples = df['study_id'].unique().tolist()
# Record statistics:
self.num_study_ids = len(self.examples)
self.num_dicom_ids = len(df['dicom_id'].unique().tolist())
self.num_subject_ids = len(df['subject_id'].unique().tolist())
def __len__(self):
return self.num_study_ids
def __getitem__(self, index):
study_id = self.examples[index]
# Get the study:
study = self.records.connect.sql(
f"""
SELECT dicom_id, study_id, subject_id, study_datetime, ViewPosition
FROM mimic_cxr
WHERE (study_id = {study_id});
"""
).df()
subject_id = study.iloc[0, study.columns.get_loc('subject_id')]
study_id = study.iloc[0, study.columns.get_loc('study_id')]
study_datetime = study['study_datetime'].max()
example_dict = {
'study_ids': study_id,
'subject_id': subject_id,
'index': index,
}
example_dict.update(self.records.return_mimic_cxr_features(study_id))
if self.ed_module:
edstays = self.records.connect.sql(
f"""
SELECT stay_id, intime, outtime
FROM edstays
WHERE (subject_id = {subject_id})
AND intime < '{study_datetime}'
AND outtime > '{study_datetime}';
"""
).df()
assert len(edstays) <= 1
stay_id = edstays.iloc[0, edstays.columns.get_loc('stay_id')] if not edstays.empty else None
self.records.clear_start_end_times()
example_dict.update(self.records.return_ed_module_features(stay_id, study_datetime))
example_dict['stay_ids'] = stay_id
if self.return_images:
example_dict['images'], example_dict['image_time_deltas'] = self.get_images(study, study_datetime)
return example_dict
def get_images(self, example, reference_time):
"""
Get the image/s for a given example.
Argument/s:
example - dataframe for the example.
reference_time - reference_time for time delta.
Returns:
The image/s for the example
"""
# Sample if over max_images_per_study. Only allowed during training:
if len(example) > self.max_images_per_study:
assert self.split == 'train'
example = example.sample(n=self.max_images_per_study, axis=0)
# Order by ViewPostion:
example['ViewPosition'] = example['ViewPosition'].astype(pd.CategoricalDtype(categories=VIEW_ORDER, ordered=True))
# Sort the DataFrame based on the categorical column
example = example.sort_values(by=['study_datetime', 'ViewPosition'])
# Load and pre-process each CXR:
images, time_deltas = [], []
for _, row in example.iterrows():
images.append(
self.load_and_preprocess_image(
row['subject_id'],
row['study_id'],
row['dicom_id'],
),
)
time_deltas.append(self.records.compute_time_delta(row['study_datetime'], reference_time, to_tensor=False))
if self.transforms is not None:
images = torch.stack(images, 0)
return images, time_deltas
def load_and_preprocess_image(self, subject_id, study_id, dicom_id):
"""
Load and preprocess an image using torchvision.transforms.v2:
https://pytorch.org/vision/stable/auto_examples/transforms/plot_transforms_getting_started.html#sphx-glr-auto-examples-transforms-plot-transforms-getting-started-py
Argument/s:
subject_id - subject identifier.
study_id - study identifier.
dicom_id - DICOM identifier.
Returns:
image - Tensor of the CXR.
"""
if self.extension == 'jpg':
if self.jpg_rocksdb_path is not None:
# Convert to bytes:
key = bytes(dicom_id, 'utf-8')
# Retrieve image:
image = bytearray(self.images_db.get(key))
image = torch.frombuffer(image, dtype=torch.uint8)
image = decode_image(image)
elif self.jpg_lmdb_path is not None:
# Convert to bytes:
key = bytes(dicom_id, 'utf-8')
# Retrieve image:
image = bytearray(self.txn.get(key))
image = torch.frombuffer(image, dtype=torch.uint8)
image = decode_image(image)
else:
image_file_path = mimic_cxr_image_path(self.dataset_dir, subject_id, study_id, dicom_id, self.extension)
image = read_image(image_file_path)
elif self.extension == 'dcm':
if self.images_rocksdb_path is not None:
key = dicom_id.encode('utf-8')
# Retrieve the serialized image shape associated with the key:
shape_bytes = self.images_db.get((self.shape_handle, key), key)
shape = struct.unpack('iii', shape_bytes)
np.frombuffer(shape_bytes, dtype=self.shape_dtype).reshape(3)
# Retrieve the serialized image data associated with the key:
image_bytes = self.images_db.get((self.image_handle, key), key)
image = np.frombuffer(image_bytes, dtype=self.image_dtype).reshape(*shape)
else:
image_file_path = mimic_cxr_image_path(self.dataset_dir, subject_id, study_id, dicom_id, self.extension)
image = load_and_preprocess_dcm_uint16(image_file_path)
# Convert to a torch tensor:
image = torch.from_numpy(image)
if self.transforms is not None:
image = self.transforms(image)
return image
if __name__ == '__main__':
import time
from tqdm import tqdm
num_samples = 20
datasets = []
datasets.append(
StudyIDEDStayIDSubset(
dataset_dir='/datasets/work/hb-mlaifsp-mm/work/archive',
mimic_iv_duckdb_path='/scratch3/nic261/database/mimic_iv_duckdb_rev_b.db',
split='train',
extension='jpg',
ed_module=False,
),
)
datasets.append(
StudyIDEDStayIDSubset(
dataset_dir='/scratch3/nic261/datasets',
mimic_iv_duckdb_path='/scratch3/nic261/database/mimic_iv_duckdb_rev_b.db',
split='train',
extension='jpg',
ed_module=False,
),
)
datasets.append(
StudyIDEDStayIDSubset(
jpg_lmdb_path='/scratch3/nic261/database/mimic_cxr_jpg_lmdb_rev_a.db',
mimic_iv_duckdb_path='/scratch3/nic261/database/mimic_iv_duckdb_rev_b.db',
split='train',
extension='jpg',
ed_module=False,
),
)
datasets.append(
StudyIDEDStayIDSubset(
jpg_rocksdb_path='/scratch3/nic261/database/mimic_cxr_jpg_rocksdb.db',
mimic_iv_duckdb_path='/scratch3/nic261/database/mimic_iv_duckdb_rev_b.db',
split='train',
extension='jpg',
ed_module=False,
)
)
assert (datasets[1][0]['images'][0] == datasets[2][0]['images'][0]).all().item()
assert (datasets[1][5]['images'][0] == datasets[2][5]['images'][0]).all().item()
for d in datasets:
start_time = time.time()
indices = torch.randperm(len(d))[:num_samples] # Get random indices.
for i in tqdm(indices):
_ = d[i]
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Elapsed time: {elapsed_time} seconds")