--- a +++ b/src/preprocessing.py @@ -0,0 +1,71 @@ +import numpy as np +import pandas as pd +import os +import click +import glob +import cv2 +import pydicom +from tqdm import tqdm +from utils import get_windowing, window_image +from joblib import delayed, Parallel + + +@click.group() +def cli(): + print("CLI") + + +windows_range = { + 'brain': [40, 80], + 'bone': [600, 2800], + 'subdual': [75, 215] +} + + +def convert_dicom_to_jpg(dicomfile, outputdir): + try: + data = pydicom.read_file(dicomfile) + image = data.pixel_array + window_center, window_width, intercept, slope = get_windowing(data) + id = dicomfile.split("/")[-1].split(".")[0] + + images = [] + for k, v in windows_range.items(): + image_windowed = window_image(image, v[0], v[1], intercept, slope) + images.append(image_windowed) + + images = np.asarray(images).transpose((1, 2, 0)) + output_image = os.path.join(outputdir, id + ".jpg") + cv2.imwrite(output_image, images) + except: + print(dicomfile) + + +@cli.command() +@click.option('--inputdir', type=str) +@click.option('--outputdir', type=str) +def extract_images( + inputdir, + outputdir, +): + os.makedirs(outputdir, exist_ok=True) + files = glob.glob(inputdir + "/*.dcm") + Parallel(n_jobs=8)(delayed(convert_dicom_to_jpg)(file, outputdir) for file in tqdm(files, total=len(files))) + + +def split_by_patient( + train_csv, + train_meta_csv, + n_folds, + outdir +): + os.makedirs(outdir, exist_ok=True) + train_df = pd.read_csv(train_csv) + train_meta_df = pd.read_csv(train_meta_csv) + train_meta_df['ID'] = train_meta_df['ID'].apply(lambda x: "_".join(x.split("_")[:2])) + train_meta_df = train_meta_df[['ID', 'PatientID']] + + + +if __name__ == '__main__': + cli()