feat: add data preprocessing pipeline

prepare_data.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from typing import Dict, List1
+
+import argparse
+import json
+from functools import partial
+import pathlib
+import shutil
+import re
+
+from tqdm import tqdm
+from PIL import Image
+import pandas as pd
+
+ImageCaptionMap = Dict[str, Dict[str, str]]
+
+def _get_image_path(row: pd.Series, root_dir: str = '.') -> str:
+    path = [
+        root_dir,
+        'files',
+        f'p{row.subject_id}'[:3],
+        f'p{row.subject_id}',
+        f's{row.study_id}',
+        f'{row.dicom_id}.jpg'
+    ]
+    
+    return '/'.join(path)
+
+def _prepare_dataframe(
+        captions: pd.DataFrame,
+        metadata: pd.DataFrame,
+        row: pd.Series
+) -> pd.Series:
+    if f's{row.study_id}' in captions.index:
+        row[captions.columns] = (
+            captions
+            .loc[f's{row.study_id}']
+            .apply(lambda text: (
+                re.sub('_+', '_', text)
+                .replace('\n', ' ')
+                .lower().rstrip('.')
+            ))
+        )
+
+    if row.dicom_id in metadata.index:
+        row['view_position'] = metadata.loc[row.dicom_id, 'ViewPosition']
+    
+    return row
+
+def copy_image(
+        row: pd.Series,
+        target_path: pathlib.Path,
+        split: str,
+        size: int = 224
+) -> str:
+    target_img_path = target_path / split / f'{row.dicom_id}.jpg'
+    target_img_path = str(target_img_path.resolve())
+    
+    img = Image.open(row.path)
+    img = img.resize((size, size))
+    img.save(target_img_path)
+    
+    return target_img_path
+
+def generate_dataset(
+        root_dir: pathlib.Path,
+        target_dir: pathlib.Path,
+        split: str = 'validate'
+) -> ImageCaptionMap:
+    meta_dir = root_dir / 'metadata'
+
+    metadata = pd.read_csv(meta_dir / 'mimic-cxr-2.0.0-metadata.csv')
+    df_split = pd.read_csv(meta_dir / 'mimic-cxr-2.0.0-split.csv')
+    captions = pd.read_csv(meta_dir / 'mimic_cxr_sectioned.csv')
+
+    captions = captions.where(~captions.isna(), '').set_index('study')
+    metadata = metadata.set_index('dicom_id')
+    
+    if split in df_split.split.unique():
+        current_split = df_split[df_split.split == split]
+        get_abs_path = partial(_get_image_path, root_dir=str(root_dir.resolve()))
+
+        current_split['path'] = current_split.apply(get_abs_path, axis=1)
+        current_split['view_position'] = ''
+        for col in captions.columns:
+            current_split[col] = ''
+
+        preprocess_func = partial(_prepare_dataframe, captions, metadata)
+
+        df = current_split.apply(preprocess_func, axis=1)
+        
+    else:
+        raise ValueError('bad split')
+    
+    image_path_to_caption = {}
+    (target_dir / split).mkdir(exist_ok=True, parents=True)
+
+    for _, element in tqdm(df.iterrows()):
+        caption = {
+            'impression': element['impression'],
+            'findings': element['findings'],
+            'last_paragraph': element['last_paragraph'],
+            'comparison': element['comparison'],
+            'view_position': element['view_position'],
+        }
+
+        image_path = copy_image(element, target_dir, split)
+
+        image_path_to_caption[image_path] = caption
+        
+    return image_path_to_caption
+
+def dump_dataset(image_path_to_caption: ImageCaptionMap) -> List[str]:
+    lines = []
+    
+    for image_path, captions in image_path_to_caption.items():
+        lines.append(json.dumps({
+            'image_path': image_path, 
+            'caption': captions,
+        }))
+    
+    return lines
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Preprocess MIMIC-CXR dataset')
+    parser.add_argument('--data_dir', description='MIMIC-CXR path')
+    parser.add_argument('--target_dir', description='output path')
+
+    args = parser.parse_args()
+
+    data_dir = pathlib.Path(args.data_dir)
+    target_dir = pathlib.Path(args.target_dir)
+
+    for split in ['test', 'validate', 'train']:
+        image_path_to_caption = generate_dataset(data_dir, target_dir, split)
+        lines = dump_dataset(image_path_to_caption)
+
+        with open(target_dir / f'{split}_dataset.json', 'w') as f:
+            f.write('\n'.join(lines))