import numpy as np
from tqdm import tqdm
from ast import literal_eval
from sybil.datasets.nlst import NLST_Survival_Dataset
from collections import Counter
import copy
DEVICE_ID = {
"GE MEDICAL SYSTEMS": 0,
"TOSHIBA": 1,
"Philips": 2,
"SIEMENS": 3,
"Siemens Healthcare": 3, # note: same id as SIEMENS
"Vital Images, Inc.": 4,
"Hitachi Medical Corporation": 5,
"LightSpeed16": 6,
}
class MGH_Dataset(NLST_Survival_Dataset):
"""
MGH Dataset Cohort 1
"""
def create_dataset(self, split_group):
"""
Gets the dataset from the paths and labels in the json.
Arguments:
split_group(str): One of ['train'|'dev'|'test'].
Returns:
The dataset as a dictionary with img paths, label,
and additional information regarding exam or participant
"""
dataset = []
# if split probs is set, randomly assign new splits, (otherwise default is 70% train, 15% dev and 15% test)
if self.args.assign_splits:
np.random.seed(self.args.cross_val_seed)
self.assign_splits(self.metadata_json)
for mrn_row in tqdm(self.metadata_json):
pid, split, exams = mrn_row["pid"], mrn_row["split"], mrn_row["accessions"]
# pt_metadata missing
for exam_dict in exams:
studyuid = exam_dict["StudyInstanceUID"]
bridge_uid = exam_dict["bridge_uid"]
days_to_last_exam = -int(
exam_dict["diff_days"]
) # no. of days to the oldest exam (0 or a negative int)
exam_no = self.get_exam_no(days_to_last_exam, exams)
y, y_seq, y_mask, time_at_event = self.get_label(exam_dict, exams)
for series_id, series_dict in exam_dict["image_series"].items():
if self.skip_sample(series_dict, exam_dict, mrn_row, split_group):
continue
img_paths = series_dict["paths"]
img_paths = [p.replace("Data082021", "pngs") for p in img_paths]
slice_locations = series_dict["image_posn"]
series_data = series_dict["series_data"]
device = DEVICE_ID[series_data["Manufacturer"]]
sorted_img_paths, sorted_slice_locs = self.order_slices(
img_paths, slice_locations
)
sample = {
"paths": sorted_img_paths,
"slice_locations": sorted_slice_locs,
"y": int(y),
"time_at_event": time_at_event,
"y_seq": y_seq,
"y_mask": y_mask,
"exam": int(
"{}{}".format(
studyuid.replace(".", "")[-5:],
series_id.replace(".", "")[-5:],
)
), # last 5 of study id + last 5 of series id
"exam_str": "{}_{}".format(bridge_uid, exam_no),
"accession": exam_no,
"study": studyuid,
"series": series_id,
"pid": pid,
"device": device,
"lung_rads": -1
if exam_dict["lung_rads"] == np.nan
else exam_dict["lung_rads"],
"IV_contrast": exam_dict["IV_contrast"],
"lung_cancer_screening": exam_dict["lung_cancer_screening"],
"cancer_location": np.zeros(14), # mgh has no annotations
"cancer_laterality": np.zeros(
3, dtype=np.int
), # has to be int, while cancer_location has to be float
"num_original_slices": len(series_dict["paths"]),
"annotations": [],
"pixel_spacing": series_dict["pixel_spacing"]
+ [series_dict["slice_thickness"]],
"slice_thickness": self.get_slice_thickness_class(
series_dict["slice_thickness"]
),
}
if self.args.use_risk_factors:
sample["risk_factors"] = self.get_risk_factors(
exam_dict, return_dict=False
)
if self.args.use_annotations:
# mgh has no annotations, so set everything to zero / false
sample["volume_annotations"] = np.array(
[0 for _ in sample["paths"]]
)
sample["annotations"] = [
{"image_annotations": None} for path in sample["paths"]
]
dataset.append(sample)
return dataset
def skip_sample(self, series_dict, exam_dict, mrn_row, split):
if not mrn_row["split"] == split:
return True
if mrn_row["in_cohort2"]:
return True
# check if screen is localizer screen or not enough images
if self.is_localizer(series_dict["series_data"]):
return True
slice_thickness = series_dict["slice_thickness"]
# check if restricting to specific slice thicknesses
if (self.args.slice_thickness_filter is not None) and (
(slice_thickness in ["", None])
or (slice_thickness > self.args.slice_thickness_filter)
or (slice_thickness < 0)
):
return True
if series_dict["pixel_spacing"] is None:
return True
# remove where slice location doesn't change (different axis):
if len(set(series_dict["image_posn"])) < 2:
return True
if len(series_dict["paths"]) < self.args.min_num_images:
return True
return False
def get_exam_no(self, diff_days, exams):
"""Gets the index of the exam, compared to the other exams"""
sorted_days = sorted([-exam["diff_days"] for exam in exams], reverse=True)
return sorted_days.index(diff_days)
def get_label(self, exam_dict, exams):
is_cancer_cohort = exam_dict["cancer_cohort_yes_no"] == "yes"
days_to_last_followup = -exam_dict["diff_days"]
years_to_last_followup = days_to_last_followup // 365
y = 0
y_seq = np.zeros(self.args.max_followup)
if is_cancer_cohort:
days_to_cancer = -exam_dict["diff_days_exam_lung_cancer_diagnosis"]
years_to_cancer = int(days_to_cancer // 365)
y = years_to_cancer < self.args.max_followup
time_at_event = min(years_to_cancer, self.args.max_followup - 1)
y_seq[years_to_cancer:] = 1
else:
time_at_event = min(years_to_last_followup, self.args.max_followup - 1)
y_mask = np.array(
[1] * (time_at_event + 1)
+ [0] * (self.args.max_followup - (time_at_event + 1))
)
y_mask = y_mask[: self.args.max_followup]
return y, y_seq.astype("float64"), y_mask.astype("float64"), time_at_event
def get_risk_factors(self, exam_dict, return_dict=False):
risk_factors = {
"age_at_exam": exam_dict["age_at_exam"],
"pack_years": exam_dict["pack_years"],
"race": exam_dict["race"],
"sex": exam_dict["sex"],
"smoking_status": exam_dict["smoking_status"],
}
if return_dict:
return risk_factors
else:
return np.array(
[v for v in risk_factors.values() if not isinstance(v, str)]
)
def is_localizer(self, series_dict):
is_localizer = "LOCALIZER" in literal_eval(series_dict["ImageType"])
return is_localizer
@staticmethod
def set_args(args):
args.num_classes = args.max_followup
def get_summary_statement(self, dataset, split_group):
summary = "Constructed MGH CT Cancer Survival {} dataset with {} records, {} exams, {} patients, and the following class balance \n {}"
class_balance = Counter([d["y"] for d in dataset])
exams = set([d["exam"] for d in dataset])
patients = set([d["pid"] for d in dataset])
statement = summary.format(
split_group,
len(dataset),
len(exams),
len(patients),
class_balance,
)
statement += "\n" + "Censor Times: {}".format(
Counter([d["time_at_event"] for d in dataset])
)
return statement
def assign_splits(self, meta):
for idx in range(len(meta)):
meta[idx]["split"] = np.random.choice(
["train", "dev", "test"], p=self.args.split_probs
)
class MGH_Screening(NLST_Survival_Dataset):
"""
MGH Dataset Cohort 2
"""
def create_dataset(self, split_group):
"""
Gets the dataset from the paths and labels in the json.
Arguments:
split_group(str): One of ['train'|'dev'|'test'].
Returns:
The dataset as a dictionary with img paths, label,
and additional information regarding exam or participant
"""
assert not self.args.train, "Cohort 2 should not be used for training"
dataset = []
for mrn_row in tqdm(self.metadata_json):
pid, exams = mrn_row["pid"], mrn_row["accessions"]
for exam_dict in exams:
for series_id, series_dict in exam_dict["image_series"].items():
if self.skip_sample(series_dict, exam_dict, mrn_row):
continue
sample = self.get_volume_dict(
series_id, series_dict, exam_dict, mrn_row
)
if len(sample) == 0:
continue
dataset.append(sample)
return dataset
def skip_sample(self, series_dict, exam_dict, mrn_row):
# unknown cancer status
if exam_dict["Future_cancer"] == "unkown":
return True
if (exam_dict["days_before_cancer_dx"] < 0) or (
exam_dict["days_to_last_follow_up"] < 0
):
return True
# check if screen is localizer screen or not enough images
if self.is_localizer(series_dict["series_data"]):
return True
slice_thickness = series_dict["SliceThickness"]
# check if restricting to specific slice thicknesses
if (self.args.slice_thickness_filter is not None) and (
(slice_thickness in ["", None])
or (slice_thickness > self.args.slice_thickness_filter)
or (slice_thickness < 0)
):
return True
if series_dict["PixelSpacing"] is None:
return True
if len(series_dict["paths"]) < self.args.min_num_images:
return True
return False
def get_volume_dict(self, series_id, series_dict, exam_dict, mrn_row):
img_paths = series_dict["paths"]
img_paths = [
p.replace("MIT_Lung_Cancer_Screening", "screening_pngs").replace(
".dcm", ".png"
)
for p in img_paths
]
slice_locations = series_dict["slice_location"]
series_data = series_dict["series_data"]
pixel_spacing = series_dict["PixelSpacing"] + [series_dict["SliceThickness"]]
sorted_img_paths, sorted_slice_locs = self.order_slices(
img_paths, slice_locations, reverse=True
)
device = DEVICE_ID[series_data["Manufacturer"]]
studyuid = exam_dict["StudyInstanceUID"]
bridge_uid = exam_dict["bridge_uid"]
y, y_seq, y_mask, time_at_event = self.get_label(exam_dict, mrn_row)
sample = {
"paths": sorted_img_paths,
"slice_locations": sorted_slice_locs,
"y": int(y),
"time_at_event": time_at_event,
"y_seq": y_seq,
"y_mask": y_mask,
"exam": int(
"{}{}".format(
studyuid.replace(".", "")[-5:],
series_id.replace(".", "")[-5:],
)
), # last 5 of study id + last 5 of series id
"study": studyuid,
"series": series_id,
"pid": mrn_row["pid"],
"bridge_uid": bridge_uid,
"device": device,
"lung_rads": exam_dict["LR Score"],
"cancer_location": np.zeros(14), # mgh has no annotations
"cancer_laterality": np.zeros(
3, dtype=np.int
), # has to be int, while cancer_location has to be float
"num_original_slices": len(series_dict["paths"]),
"marital_status": exam_dict["marital_status"],
"religion": exam_dict["religion"],
"primary_site": exam_dict["Primary Site"],
"laterality1": exam_dict["Laterality"],
"laterality2": exam_dict["Laterality.1"],
"icdo3": exam_dict["Histo/Behavior ICD-O-3"],
"pixel_spacing": pixel_spacing,
"slice_thickness": self.get_slice_thickness_class(pixel_spacing[-1]),
}
if self.args.use_risk_factors:
sample["risk_factors"] = self.get_risk_factors(exam_dict, return_dict=False)
if self.args.use_annotations:
# mgh has no annotations, so set everything to zero / false
sample["volume_annotations"] = np.array([0 for _ in sample["paths"]])
sample["annotations"] = [
{"image_annotations": None} for path in sample["paths"]
]
return sample
def get_label(self, exam_dict, mrn_row):
is_cancer_cohort = exam_dict["Future_cancer"].lower().strip() == "yes"
days_to_cancer = exam_dict["days_before_cancer_dx"]
y = False
if (
is_cancer_cohort
and (not np.isnan(days_to_cancer))
and (days_to_cancer > -1)
):
years_to_cancer = int(days_to_cancer // 365)
y = years_to_cancer < self.args.max_followup
y_seq = np.zeros(self.args.max_followup)
if y:
time_at_event = years_to_cancer
y_seq[years_to_cancer:] = 1
else:
if is_cancer_cohort:
assert (days_to_cancer < 0) or (
years_to_cancer >= self.args.max_followup
)
time_at_event = self.args.max_followup - 1
else:
days_to_last_neg_followup = exam_dict["days_to_last_follow_up"]
years_to_last_neg_followup = int(days_to_last_neg_followup // 365)
time_at_event = min(
years_to_last_neg_followup, self.args.max_followup - 1
)
y_mask = np.array(
[1] * (time_at_event + 1)
+ [0] * (self.args.max_followup - (time_at_event + 1))
)
y_mask = y_mask[: self.args.max_followup]
return y, y_seq.astype("float64"), y_mask.astype("float64"), time_at_event
def get_risk_factors(self, exam_dict, return_dict=False):
risk_factors = {
"race": exam_dict["race"],
"pack_years": exam_dict["Packs Years"],
"age_at_exam": exam_dict["age at the exam"],
"gender": exam_dict["gender"],
"smoking_status": exam_dict["Smoking Status"],
"lung_rads": exam_dict["LR Score"],
"years_since_quit_smoking": exam_dict["Year Since Last Smoked"],
}
if return_dict:
return risk_factors
else:
return np.array(
[v for v in risk_factors.values() if not isinstance(v, str)]
)
def is_localizer(self, series_dict):
is_localizer = "LOCALIZER" in literal_eval(series_dict["ImageType"])
return is_localizer
@staticmethod
def set_args(args):
args.num_classes = args.max_followup
def get_summary_statement(self, dataset, split_group):
summary = "Constructed MGH CT Cancer Survival {} dataset with {} records, {} exams, {} patients, and the following class balance \n {}"
class_balance = Counter([d["y"] for d in dataset])
exams = set([d["exam"] for d in dataset])
patients = set([d["pid"] for d in dataset])
statement = summary.format(
split_group,
len(dataset),
len(exams),
len(patients),
class_balance,
)
statement += "\n" + "Censor Times: {}".format(
Counter([d["time_at_event"] for d in dataset])
)
return statement
def assign_splits(self, meta):
for idx in range(len(meta)):
meta[idx]["split"] = np.random.choice(
["train", "dev", "test"], p=self.args.split_probs
)
Datasets
Models
