Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update some section for easier readability. #25

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open

Update some section for easier readability. #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
022a93a
Update convert_ckpt_to_pt.py
Geeks-Sid May 10, 2023
d6ef3da
Update rex_o4a_csv
Geeks-Sid May 10, 2023
8add314
Update rex_sin_csv
Geeks-Sid May 10, 2023
2fc6c55
Update rex_mul_csv
Geeks-Sid May 10, 2023
b6c8468
Update rex_bids_csv
Geeks-Sid May 10, 2023
7e90546
Update csv_creator_adv.py
Geeks-Sid May 10, 2023
a3c88f1
Update optimizers.py
Geeks-Sid May 10, 2023
57b2dca
Update utils_test.py
Geeks-Sid May 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 9 additions & 7 deletions BrainMaGe/utils/convert_ckpt_to_pt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,16 @@
ckpt_file = os.path.abspath(args.input)
pt_file = os.path.abspath(args.output)

print("Attempting to load file : ", ckpt_file)
weight_load = torch.load(ckpt_file)
print("Load Successful! Converting file.")
print("Attempting to load file:", ckpt_file)
with open(ckpt_file, "rb") as f:
weight_load = torch.load(f, map_location=torch.device("cpu"))
print("Load successful! Converting file.")
new_state_dict = {}
for key in weight_load["state_dict"].keys():
for key, value in weight_load["state_dict"].items():
new_key = key[6:]
new_state_dict[new_key] = weight_load["state_dict"][key]
new_state_dict[new_key] = value
model_state_dict = {"model_state_dict": new_state_dict}
print("Conversion successful!")
torch.save(model_state_dict, pt_file)
print("File saved successfully at :", pt_file)
with open(pt_file, "wb") as f:
torch.save(model_state_dict, f)
print("File saved successfully at:", pt_file)
301 changes: 139 additions & 162 deletions BrainMaGe/utils/csv_creator_adv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,233 +14,210 @@
from bids import BIDSLayout


def rex_o4a_csv(folder_path, to_save, ftype, modalities):
"""[CSV generation for OneForAll]
[This function is used to generate a csv for OneForAll mode and creates a
csv]
Arguments:
folder_path {[string]} -- [Takes the folder to see where to look for
the different modaliies]
to_save {[string]} -- [Takes the folder as a string to save the csv]
ftype {[string]} -- [Are you trying to save train, validation or test,
if file type is set to test, it does not look for
ground truths]
modalities {[string]} -- [usually a string which looks like this
: ['t1', 't2', 't1ce']]
def rex_o4a_csv(folder_path, save_folder, file_type, modalities):
"""
CSV generation for OneForAll.
This function generates a csv for OneForAll mode and creates a csv.

Args:
folder_path (str): The folder to see where to look for the different modalities
save_folder (str): The folder to save the csv
file_type (str): train, validation or test. If file_type is set to test, it does not look for ground truths.
modalities (list of str): The modalities to include in the csv.
"""
modalities = modalities[1:-1]
modalities = re.findall("[^, ']+", modalities)
if not modalities:
print(
"Could not find modalities! Are you sure you have put in \
something in the modalities field?"
)
sys.exit(0)
if ftype == "test":
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,Image_Path\n")
raise ValueError("Could not find modalities! Are you sure you have put in something in the modalities field?")

if file_type == "test":
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,Image_Path\n")
else:
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,gt_path,Image_path\n")
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,gt_path,Image_path\n")

folders = os.listdir(folder_path)
for folder in folders:
for modality in modalities:
csv_file.write(folder + "_" + modality + ",")
if ftype != "test":
ground_truth = glob.glob(
os.path.join(folder_path, folder, "*mask.nii.gz")
)[0]
csv_file.write(ground_truth)

if file_type != "test":
ground_truth_path = glob.glob(os.path.join(folder_path, folder, "*mask.nii.gz"))[0]
csv_file.write(ground_truth_path)
csv_file.write(",")
img = glob.glob(
os.path.join(folder_path, folder, "*" + modality + ".nii.gz")
)[0]
csv_file.write(img)

image_path = glob.glob(os.path.join(folder_path, folder, "*" + modality + ".nii.gz"))[0]
csv_file.write(image_path)
csv_file.write("\n")
csv_file.close()

print("CSV file saved successfully at:", csv_path)


def rex_sin_csv(folder_path, to_save, ftype, modalities):
"""[CSV generation for Single Modalities]
[This function is used to generate a csv for Single mode and creates a csv]
Arguments:
folder_path {[string]} -- [Takes the folder to see where to look for
the different modaliies]
to_save {[string]} -- [Takes the folder as a string to save the csv]
ftype {[string]} -- [Are you trying to save train, validation or test,
if file type is set to test, it does not look for
ground truths]
modalities {[string]} -- [usually a string which looks like this
: ['t1']]
def rex_sin_csv(folder_path, save_folder, file_type, modalities):
"""
CSV generation for Single Modalities.
This function generates a csv for Single mode and creates a csv.

Args:
folder_path (str): The folder to see where to look for the different modalities
save_folder (str): The folder to save the csv
file_type (str): train, validation or test. If file_type is set to test, it does not look for ground truths.
modalities (list of str): The modalities to include in the csv.
"""
modalities = modalities[1:-1]
modalities = re.findall("[^, ']+", modalities)
if len(modalities) > 1:
print("Found more than one modality, exiting!")
sys.exit(0)
raise ValueError("Found more than one modality, exiting!")
if not modalities:
print(
"Could not find modalities! Are you sure you have put in \
something in the modalities field?"
)
sys.exit(0)
if ftype == "test":
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,")
raise ValueError("Could not find modalities! Are you sure you have put in something in the modalities field?")

if file_type == "test":
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,")
else:
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,gt_path,")
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,gt_path,")

modality = modalities[0]
csv_file.write(modality + "_path\n")
folders = os.listdir(folder_path)
for folder in folders:
csv_file.write(folder)
csv_file.write(",")
if ftype != "test":
ground_truth = glob.glob(os.path.join(folder_path, folder, "*mask.nii.gz"))[
0
]
csv_file.write(ground_truth)
if file_type != "test":
ground_truth_path = glob.glob(os.path.join(folder_path, folder, "*mask.nii.gz"))[0]
csv_file.write(ground_truth_path)
csv_file.write(",")
img = glob.glob(os.path.join(folder_path, folder, "*" + modality + ".nii.gz"))[
0
]
csv_file.write(img)

image_path = glob.glob(os.path.join(folder_path, folder, "*" + modality + ".nii.gz"))[0]
csv_file.write(image_path)
csv_file.write("\n")
csv_file.close()

print("CSV file saved successfully at:", csv_path)


def rex_mul_csv(folder_path, to_save, ftype, modalities):
"""[CSV generation for Multi Modalities]
[This function is used to generate a csv for multi mode and creates a csv]
Arguments:
folder_path {[string]} -- [Takes the folder to see where to look for
the different modaliies]
to_save {[string]} -- [Takes the folder as a string to save the csv]
ftype {[string]} -- [Are you trying to save train, validation or test,
if file type is set to test, it does not look for
ground truths]
modalities {[string]} -- [usually a string which looks like this
: ['t1']]
def rex_mul_csv(folder_path, save_folder, file_type, modalities):
"""
CSV generation for Multi Modalities.
This function generates a csv for multi mode and creates a csv.

Args:
folder_path (str): The folder to see where to look for the different modalities
save_folder (str): The folder to save the csv
file_type (str): train, validation or test. If file_type is set to test, it does not look for ground truths.
modalities (list of str): The modalities to include in the csv.
"""
modalities = modalities[1:-1]
modalities = re.findall("[^, ']+", modalities)
if not modalities:
print(
"Could not find modalities! Are you sure you have put in \
something in the modalities field?"
)
sys.exit(0)
if ftype == "test":
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,")
raise ValueError("Could not find modalities! Are you sure you have put in something in the modalities field?")

if file_type == "test":
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,")
else:
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,gt_path,")
for modality in modalities[:-1]:
csv_file.write(modality + "_path,")
modality = modalities[-1]
csv_file.write(modality + "_path\n")
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,gt_path,")

csv_file.write(",".join([f"{modality}_path" for modality in modalities]))
csv_file.write("\n")

folders = os.listdir(folder_path)
for folder in folders:
csv_file.write(folder)
csv_file.write(",")
if ftype != "test":
ground_truth = glob.glob(os.path.join(folder_path, folder, "*mask.nii.gz"))[
0
]
csv_file.write(ground_truth)

if file_type != "test":
ground_truth_path = glob.glob(os.path.join(folder_path, folder, "*mask.nii.gz"))[0]
csv_file.write(ground_truth_path)
csv_file.write(",")

for modality in modalities[:-1]:
img = glob.glob(
os.path.join(folder_path, folder, "*" + modality + ".nii.gz")
)[0]
csv_file.write(img)
image_path = glob.glob(os.path.join(folder_path, folder, "*" + modality + ".nii.gz"))[0]
csv_file.write(image_path)
csv_file.write(",")
modality = modalities[-1]
img = glob.glob(os.path.join(folder_path, folder, "*" + modality + ".nii.gz"))[
0
]
csv_file.write(img)

image_path = glob.glob(os.path.join(folder_path, folder, "*" + modalities[-1] + ".nii.gz"))[0]
csv_file.write(image_path)
csv_file.write("\n")
csv_file.close()

print("CSV file saved successfully at:", csv_path)


def rex_bids_csv(folder_path, to_save, ftype):
"""[CSV generation for BIDS datasets]
[This function is used to generate a csv for BIDS datasets]
Arguments:
folder_path {[string]} -- [Takes the folder to see where to look for
the different modaliies]
to_save {[string]} -- [Takes the folder as a string to save the csv]
ftype {[string]} -- [Are you trying to save train, validation or test,
if file type is set to test, it does not look for
ground truths]
def rex_bids_csv(folder_path, save_folder, file_type):
"""
CSV generation for BIDS datasets.
This function generates a csv for BIDS datasets.

Args:
folder_path (str): The folder to see where to look for the different modalities
save_folder (str): The folder to save the csv
file_type (str): train, validation or test. If file_type is set to test, it does not look for ground truths.
"""
if ftype == "test":
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,")
if file_type == "test":
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,")
else:
csv_file = open(os.path.join(to_save, ftype + ".csv"), "w+")
csv_file.write("ID,gt_path,")
# load BIDS dataset into memory
csv_path = os.path.join(save_folder, file_type + ".csv")
with open(csv_path, "w+") as csv_file:
csv_file.write("ID,gt_path,")

layout = BIDSLayout(folder_path)
bids_df = layout.to_df()
bids_modality_df = {
"t1": bids_df[bids_df["suffix"] == "T1w"],
"t2": bids_df[bids_df["suffix"] == "T2w"],
"flair": bids_df[bids_df["suffix"] == "FLAIR"],
"t1ce": bids_df[bids_df["suffix"] == "T1CE"],
}
# check what modalities the dataset contains
modalities = []
for modality, df in bids_modality_df.items():
if not df.empty:
modalities.append(modality)
# write headers for those modalities
for modality in modalities[:-1]:
csv_file.write(modality + "_path,")
modality = modalities[-1]
csv_file.write(modality + "_path\n")
# write image paths for each subject
modalities = ["t1", "t2", "flair", "t1ce"]
modalities = [modality for modality in modalities if layout.get(suffix=f"{modality}w")]

csv_file.write(",".join([f"{modality}_path" for modality in modalities]))
csv_file.write("\n")

for sub in layout.get_subjects():
csv_file.write(sub)
csv_file.write(",")
if ftype != "test":
ground_truth = glob.glob(os.path.join(folder_path, sub, "*mask.nii.gz"))[0]
csv_file.write(ground_truth)

if file_type != "test":
ground_truth_path = layout.get(subject=sub, suffix="mask")[0].filename
csv_file.write(ground_truth_path)
csv_file.write(",")

for modality in modalities[:-1]:
img = bids_modality_df[modality][bids_df["subject"] == sub].path.values
csv_file.write(img[0])
image_path = layout.get(subject=sub, suffix=f"{modality}w")[0].filename
csv_file.write(image_path)
csv_file.write(",")
modality = modalities[-1]
img = bids_modality_df[modality][bids_df["subject"] == sub].path.values
csv_file.write(img[0])
image_path = layout.get(subject=sub, suffix=f"{modalities[-1]}w")[0].filename
csv_file.write(image_path)
csv_file.write("\n")
csv_file.close()

print("CSV file saved successfully at:", csv_path)


def generate_csv(folder_path, to_save, mode, ftype, modalities):
"""[Function to generate CSV]
[This function takes a look at the data directory and the modes and
generates a csv]
[This function takes a look at the data directory and the modes and generates a csv]
Arguments:
folder_path {[strin]} -- [description]
to_save {[strin]} -- [description]
mode {[string]} -- [description]
ftype {[string]} -- [description]
modalities {[string]} -- [description]
"""
print("Generating ", ftype, ".csv", sep="")
if mode.lower() == "ma":
rex_o4a_csv(folder_path, to_save, ftype, modalities)
elif mode.lower() == "single":
rex_sin_csv(folder_path, to_save, ftype, modalities)
elif mode.lower() == "multi":
rex_mul_csv(folder_path, to_save, ftype, modalities)
elif mode.lower() == "bids":
rex_bids_csv(folder_path, to_save, ftype)
print(f"Generating {ftype}.csv")
modes = {
"ma": rex_o4a_csv,
"single": rex_sin_csv,
"multi": rex_mul_csv,
"bids": rex_bids_csv
}
func = modes.get(mode.lower())
if func:
func(folder_path, to_save, ftype, modalities)
else:
print("Sorry, this mode is not supported")
sys.exit(0)
Loading