Upload model

dataset.py

@@ -1,9 +1,10 @@
 import os
 
+import lmdb
 import pandas as pd
 import torch
 from torch.utils.data import Dataset
-from torchvision.io import read_image
+from torchvision.io import decode_image, read_image
 
 # 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']
@@ -25,7 +26,8 @@ class StudyIDEDStayIDSubset(Dataset):
         self, 
         split, 
         records,
-        dataset_dir=None, 
+        mimic_cxr_jpg_lmdb_path=None,
+        mimic_cxr_dir=None, 
         max_images_per_study=None,
         transforms=None, 
         images=True,
@@ -39,8 +41,9 @@ class StudyIDEDStayIDSubset(Dataset):
         """
         Argument/s:
             split - 'train', 'validate', or 'test'.
-            dataset_dir - Dataset directory.
             records - MIMIC-CXR & MIMIC-IV-ED records class instance.
+            mimic_cxr_jpg_lmdb_path - JPG database for MIMIC-CXR-JPG.
+            mimic_cxr_dir - Path to the MIMIC-CXR directory containing the patient study subdirectories with the JPG or DCM images.
             max_images_per_study - the maximum number of images per study.
             transforms - torchvision transformations.
             colour_space - PIL target colour space.
@@ -54,7 +57,8 @@ class StudyIDEDStayIDSubset(Dataset):
         """
         super(StudyIDEDStayIDSubset, self).__init__()
         self.split = split
-        self.dataset_dir = dataset_dir
+        self.mimic_cxr_jpg_lmdb_path = mimic_cxr_jpg_lmdb_path
+        self.mimic_cxr_dir = mimic_cxr_dir
         self.records = records
         self.max_images_per_study = max_images_per_study
         self.transforms = transforms
@@ -68,15 +72,16 @@ class StudyIDEDStayIDSubset(Dataset):
         # 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'
+        assert self.extension == 'jpg' or self.extension == 'dcm', '"extension" can only be either "jpg" or "dcm".'
+        assert (mimic_cxr_jpg_lmdb_path is None) != (mimic_cxr_dir is None), 'Either "mimic_cxr_jpg_lmdb_path" or "mimic_cxr_dir" can be set.'
 
-        if self.dataset_dir is not None:
+        if self.mimic_cxr_dir is not None and self.mimic_cxr_jpg_lmdb_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')
+                if 'physionet.org/files/mimic-cxr-jpg/2.0.0/files' not in self.mimic_cxr_dir:
+                    self.mimic_cxr_dir = os.path.join(self.mimic_cxr_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')
+                if 'physionet.org/files/mimic-cxr/2.0.0/files' not in self.mimic_cxr_dir:
+                    self.mimic_cxr_dir = os.path.join(self.mimic_cxr_dir, 'physionet.org/files/mimic-cxr/2.0.0/files')
 
         query = f"""
         SELECT {columns}
@@ -108,6 +113,18 @@ class StudyIDEDStayIDSubset(Dataset):
         self.num_dicom_ids = len(df['dicom_id'].unique().tolist())
         self.num_subject_ids = len(df['subject_id'].unique().tolist())
 
+        # Prepare the LMDB .jpg database:
+        if self.mimic_cxr_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.mimic_cxr_jpg_lmdb_path, map_size=map_size, lock=False, readonly=True)
+            self.txn = self.env.begin(write=False)
+
     def __len__(self):
         return self.num_study_ids
 
@@ -212,9 +229,20 @@ class StudyIDEDStayIDSubset(Dataset):
         """
 
         if self.extension == 'jpg':
-
-            image_file_path = mimic_cxr_image_path(self.dataset_dir, subject_id, study_id, dicom_id, self.extension)
-            image = read_image(image_file_path)
+            
+            if self.mimic_cxr_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.mimic_cxr_dir, subject_id, study_id, dicom_id, self.extension)
+                image = read_image(image_file_path)
 
         elif self.extension == 'dcm':
             raise NotImplementedError

lmdb_jpg.py

@@ -0,0 +1,69 @@
+import multiprocessing
+
+import duckdb
+import lmdb
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from .dataset import mimic_cxr_image_path
+
+
+class JPGDataset(Dataset):
+    def __init__(self, df, jpg_path):
+        self.df = df
+        self.jpg_path = jpg_path
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        
+        row = self.df.iloc[idx]
+        
+        jpg_path = mimic_cxr_image_path(self.jpg_path, row['subject_id'], row['study_id'], row['dicom_id'], 'jpg')
+       
+        # Convert key to bytes:
+        key = bytes(row['dicom_id'], 'utf-8')
+
+        # Read the .jpg file as bytes:
+        with open(jpg_path, 'rb') as f:
+            image = f.read()
+
+        return {
+            'keys': key,
+            'images': image,
+        }
+
+def prepare_mimic_cxr_jpg_lmdb(mimic_iv_duckdb_path, mimic_cxr_jpg_path, mimic_cxr_jpg_lmdb_path, map_size_tb, num_workers=None):
+    
+    num_workers = num_workers if num_workers is not None else multiprocessing.cpu_count()
+
+    connect = duckdb.connect(mimic_iv_duckdb_path, read_only=True)
+    df = connect.sql("SELECT DISTINCT ON(dicom_id) subject_id, study_id, dicom_id FROM mimic_cxr").df()
+    connect.close()
+
+    # Map size:
+    map_size = int(map_size_tb * (1024 ** 4))
+    assert isinstance(map_size, int)
+
+    print(f'Map size: {map_size}')
+    
+    dataset = JPGDataset(df, mimic_cxr_jpg_path)
+    dataloader = DataLoader(
+        dataset, 
+        batch_size=num_workers, 
+        shuffle=False, 
+        num_workers=num_workers, 
+        prefetch_factor=1, 
+        collate_fn=lambda x: x,
+    )
+    
+    env = lmdb.open(mimic_cxr_jpg_lmdb_path, map_size=map_size, readonly=False)
+    for batch in tqdm(dataloader):
+        for i in batch:
+            with env.begin(write=True) as txn:
+                value = txn.get(b'image_keys')
+                if value is None:
+                    txn.put(i['keys'], i['images'])
+            env.sync()
+    env.close()

modelling_cxrmate_ed.py

@@ -21,6 +21,7 @@ from transformers.utils import logging
 
 from .create_section_files import create_section_files
 from .dataset import StudyIDEDStayIDSubset
+from .lmdb_jpg import prepare_mimic_cxr_jpg_lmdb
 from .modelling_uniformer import MultiUniFormerWithProjectionHead
 from .records import EDCXRSubjectRecords
 from .tables import ed_module_tables, mimic_cxr_tables
@@ -917,11 +918,14 @@ class MIMICIVEDCXRMultimodalModel(VisionEncoderDecoderModel):
         return position_ids
     
     @staticmethod
-    def prepare_data(physionet_dir, database_path, dataset_dir=None):
+    def prepare_data(physionet_dir, database_dir):
         
-        dataset_dir = physionet_dir if dataset_dir is None else dataset_dir
-        
-        sectioned_dir = os.path.join(dataset_dir, 'mimic_cxr_sectioned')
+        Path(database_dir).mkdir(parents=True, exist_ok=True)
+       
+        mimic_iv_duckdb_path = os.path.join(database_dir, 'mimic_iv_duckdb.db')
+        mimic_cxr_jpg_lmdb_path = os.path.join(database_dir, 'mimic_cxr_jpg_lmdb.db')
+
+        sectioned_dir = os.path.join(database_dir, 'mimic_cxr_sectioned')
 
         mimic_cxr_sectioned_path = os.path.join(sectioned_dir, 'mimic_cxr_sectioned.csv')
         if not os.path.exists(mimic_cxr_sectioned_path):
@@ -947,9 +951,9 @@ class MIMICIVEDCXRMultimodalModel(VisionEncoderDecoderModel):
                 no_split=True,
             )
             
-        if not os.path.exists(database_path):
+        if not os.path.exists(mimic_iv_duckdb_path):
             
-            connect = duckdb.connect(database_path)
+            connect = duckdb.connect(mimic_iv_duckdb_path)
 
             csv_paths = []         
             csv_paths.append(glob(os.path.join(physionet_dir, 'mimic-iv-ed', '*', 'ed', 'edstays.csv.gz'))[0])
@@ -982,14 +986,16 @@ class MIMICIVEDCXRMultimodalModel(VisionEncoderDecoderModel):
             # MIMIC-CXR report sections:
             print(f'Copying mimic_cxr_sectioned into database...')  
             connect.sql(f"CREATE OR REPLACE TABLE mimic_cxr_sectioned AS FROM '{mimic_cxr_sectioned_path}';")   
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column0 TO study;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column1 TO impression;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column2 TO findings;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column3 TO indication;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column4 TO history;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column5 TO last_paragraph;")
-            connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column6 TO comparison;")
-            connect.sql("DELETE FROM mimic_cxr_sectioned WHERE study='study';")
+            columns = list(connect.sql('FROM mimic_cxr_sectioned LIMIT 1').df().columns)
+            if 'column0' in columns:  # If the column headers are not read correctly:
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column0 TO study;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column1 TO impression;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column2 TO findings;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column3 TO indication;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column4 TO history;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column5 TO last_paragraph;")
+                connect.sql("ALTER TABLE mimic_cxr_sectioned RENAME COLUMN column6 TO comparison;")
+                connect.sql("DELETE FROM mimic_cxr_sectioned WHERE study='study';")
 
             splits = connect.sql("FROM mimic_cxr_2_0_0_split").df()
             reports = connect.sql("FROM mimic_cxr_sectioned").df()
@@ -1065,6 +1071,7 @@ class MIMICIVEDCXRMultimodalModel(VisionEncoderDecoderModel):
             df = df.sort_values(by='study_datetime', ascending=False)
             df = df.groupby('study_id').first().reset_index()
 
+            print('Searching for studies associated with an ED stay...')
             for _, row in tqdm(df.iterrows(), total=df.shape[0]):
                 edstays = connect.sql(
                     f"""
@@ -1109,21 +1116,39 @@ class MIMICIVEDCXRMultimodalModel(VisionEncoderDecoderModel):
                 df = pd.DataFrame(v)
                 df = df.drop_duplicates(subset=['study_id', 'stay_id'])
                 connect.sql(f"CREATE TABLE {k}_study_ids AS SELECT * FROM df")
+            
+            connect.close()
+
+        if not os.path.exists(mimic_cxr_jpg_lmdb_path):
+            print('Preparing MIMIC-CXR-JPG LMDB database...')
+            pattern = os.path.join(physionet_dir, 'mimic-cxr-jpg', '*', 'files')
+            mimic_cxr_jpg_dir = glob(pattern)
+            assert len(mimic_cxr_jpg_dir), f'Multiple directories matched the pattern {pattern}: {mimic_cxr_jpg_dir}. Only one is required.'
+            prepare_mimic_cxr_jpg_lmdb(
+                mimic_iv_duckdb_path=mimic_iv_duckdb_path, 
+                mimic_cxr_jpg_dir=mimic_cxr_jpg_dir[0], 
+                mimic_cxr_jpg_lmdb_path=mimic_cxr_jpg_lmdb_path, 
+                map_size_tb=0.65
+            )
 
     @staticmethod                
-    def get_dataset(split, transforms, database_path, mimic_cxr_jpg_dir, max_images_per_study=5, records=None):
+    def get_dataset(split, transforms, database_dir, max_images_per_study=5, mimic_cxr_jpg_dir=None, records=None):
+        
+        mimic_iv_duckdb_path = os.path.join(database_dir, 'mimic_iv_duckdb.db')
+        mimic_cxr_jpg_lmdb_path = os.path.join(database_dir, 'mimic_cxr_jpg_lmdb.db') if mimic_cxr_jpg_dir is None else None
         
         if records is None:
             
             # This is the setup for CXRs + all effective inputs - medicine reconciliation:
-            records = EDCXRSubjectRecords(database_path=database_path, time_delta_map=lambda x: 1 / math.sqrt(x + 1)) 
+            records = EDCXRSubjectRecords(database_path=mimic_iv_duckdb_path, time_delta_map=lambda x: 1 / math.sqrt(x + 1)) 
             
             records.ed_module_tables = {k: records.ed_module_tables[k] for k in ['edstays', 'triage', 'vitalsign']}
             records.mimic_cxr_tables = {k: records.mimic_cxr_tables[k] for k in ['mimic_cxr_sectioned']}
             records.mimic_cxr_tables['mimic_cxr_sectioned'].text_columns = ['indication', 'history']
         
         dataset = StudyIDEDStayIDSubset(
-                dataset_dir=mimic_cxr_jpg_dir,
+                mimic_cxr_jpg_lmdb_path=mimic_cxr_jpg_lmdb_path,
+                mimic_cxr_dir=mimic_cxr_jpg_dir,
                 transforms=transforms,
                 split=split,
                 max_images_per_study=max_images_per_study,