Skip to content

Commit

Permalink
Error handling for predict
Browse files Browse the repository at this point in the history
  • Loading branch information
@aecelaya
aecelaya committed Apr 24, 2024
1 parent bc4a3a4 commit 06bc78b
Show file tree
Hide file tree
Showing 3 changed files with 119 additions and 112 deletions.
71 changes: 34 additions & 37 deletions Dockerfile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,44 +2,41 @@ ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:23.06-py3
FROM ${FROM_IMAGE_NAME}

# Set environment variables
ENV OMP_NUM_THREADS=2
ENV OMPI_MCA_coll_hcoll_enable 0
ENV HCOLL_ENABLE_MCAST 0
#ENV OMP_NUM_THREADS=2
#ENV OMPI_MCA_coll_hcoll_enable 0
#ENV HCOLL_ENABLE_MCAST 0
ENV DEBIAN_FRONTEND noninteractive

# Install dependencies
RUN apt-get update -y --fix-missing \
&& apt-get install -y cmake git \

# Install ANTs
RUN mkdir /opt/ants \
&& cd /opt/ants \
&& git clone https://github.com/ANTsX/ANTs.git \
&& cd /opt/ants/ANTs \
&& git checkout v2.5.0 \
&& cd /opt/ants \
&& mkdir build install \
&& cd /opt/ants/build \
&& cmake -DCMAKE_INSTALL_PREFIX=/opt/ants/install ../ANTs 2>&1 | tee cmake.log \
&& make -j 4 2>&1 | tee build.log \
&& cd /opt/ants/build/ANTS-build \
&& make install 2>&1 | tee install.log

# Install c3d
RUN mkdir /opt/c3d \
&& cd /opt/c3d/ \
&& wget https://downloads.sourceforge.net/project/c3d/c3d/Nightly/c3d-nightly-Linux-x86_64.tar.gz \
&& tar -xvf c3d-nightly-Linux-x86_64.tar.gz \
&& cp c3d-1.1.0-Linux-x86_64/bin/c?d /usr/local/bin/

# Create app work dir
ENV PATH /opt/ants/install/bin:$PATH
# Install mist
RUN pip install --upgrade pip \
&& pip install --upgrade --no-cache-dir mist-medical

# Install dependencies
#RUN apt-get update -y --fix-missing \
# && apt-get install -y cmake git
#
## Install ANTs
#RUN mkdir /opt/ants \
# && cd /opt/ants \
# && git clone https://github.com/ANTsX/ANTs.git \
# && cd /opt/ants/ANTs \
# && git checkout v2.5.0 \
# && cd /opt/ants \
# && mkdir build install \
# && cd /opt/ants/build \
# && cmake -DCMAKE_INSTALL_PREFIX=/opt/ants/install ../ANTs 2>&1 | tee cmake.log \
# && make -j 4 2>&1 | tee build.log \
# && cd /opt/ants/build/ANTS-build \
# && make install 2>&1 | tee install.log \
#
## Install c3d
#RUN mkdir /opt/c3d \
# && cd /opt/c3d/ \
# && wget https://downloads.sourceforge.net/project/c3d/c3d/Nightly/c3d-nightly-Linux-x86_64.tar.gz \
# && tar -xvf c3d-nightly-Linux-x86_64.tar.gz \
# && cp c3d-1.1.0-Linux-x86_64/bin/c3d /usr/local/bin/ \

# Create working directory
#ENV PATH /opt/ants/install/bin:$PATH
RUN mkdir /app

WORKDIR /app

# Install MIST
RUN git clone https://github.com/aecelaya/MIST.git \
&& cd MIST \
&& pip install -e .
WORKDIR /app
158 changes: 84 additions & 74 deletions mist/inference/main_inference.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -215,81 +215,91 @@ def test_time_inference(df,
# Set up rich progress bar
testing_progress = get_progress_bar("Testing")

messages = list()

# Run prediction on all samples and compute metrics
with testing_progress as pb:
for ii in pb.track(range(len(df))):
patient = df.iloc[ii].to_dict()

# Create individual folders for each prediction if output_std is enabled
if output_std:
output_std_dest = os.path.join(dest, str(patient['id']))
create_empty_dir(output_std_dest)
else:
output_std_dest = dest

if "mask" in df.columns and "fold" in df.columns:
image_list = list(patient.values())[3:]
elif "mask" in df.columns or "fold" in df.columns:
image_list = list(patient.values())[2:]
else:
image_list = list(patient.values())[1:]

og_ants_img = ants.image_read(image_list[0])

if no_preprocess:
torch_img, _, fg_bbox, _ = convert_nifti_to_numpy(image_list, None)
else:
torch_img, _, fg_bbox, _ = preprocess_example(config, image_list, None, False, None)

# Make image channels first and add batch dimension
torch_img = np.transpose(torch_img, axes=(3, 0, 1, 2))
torch_img = np.expand_dims(torch_img, axis=0)

torch_img = torch.Tensor(torch_img.copy()).to(torch.float32)
torch_img = torch_img.to("cuda")

prediction, std_images = predict_single_example(torch_img,
og_ants_img,
config,
models,
overlap,
blend_mode,
tta,
output_std,
fg_bbox)

# Apply postprocessing if required
transforms = ["remove_small_objects", "top_k_cc", "fill_holes"]
for transform in transforms:
if len(config[transform]) > 0:
for i in range(len(config[transform])):
if transform == "remove_small_objects":
transform_kwargs = {"small_object_threshold": config[transform][i][1]}
if transform == "top_k_cc":
transform_kwargs = {"morph_cleanup": config[transform][i][1],
"morph_cleanup_iterations": config[transform][i][2],
"top_k": config[transform][i][3]}
if transform == "fill_holes":
transform_kwargs = {"fill_label": config[transform][i][1]}

prediction = apply_transform(prediction,
transform,
config["labels"],
config[transform][i][0],
transform_kwargs)

# Write prediction mask to nifti file and save to disk
prediction_filename = '{}.nii.gz'.format(str(patient['id']))
output = os.path.join(output_std_dest, prediction_filename)
ants.image_write(prediction, output)

# Write standard deviation image(s) to nifti file and save to disk (only for foreground labels)
if output_std:
for i in range(len(std_images)):
if config["labels"][i] > 0:
std_image_filename = '{}_std_{}.nii.gz'.format(patient['id'], config['labels'][i])
output = os.path.join(output_std_dest, std_image_filename)
ants.image_write(std_images[i], output)

# Clean up
gc.collect()
try:
# Create individual folders for each prediction if output_std is enabled
if output_std:
output_std_dest = os.path.join(dest, str(patient['id']))
create_empty_dir(output_std_dest)
else:
output_std_dest = dest

if "mask" in df.columns and "fold" in df.columns:
image_list = list(patient.values())[3:]
elif "mask" in df.columns or "fold" in df.columns:
image_list = list(patient.values())[2:]
else:
image_list = list(patient.values())[1:]

og_ants_img = ants.image_read(image_list[0])

if no_preprocess:
torch_img, _, fg_bbox, _ = convert_nifti_to_numpy(image_list, None)
else:
torch_img, _, fg_bbox, _ = preprocess_example(config, image_list, None, False, None)

# Make image channels first and add batch dimension
torch_img = np.transpose(torch_img, axes=(3, 0, 1, 2))
torch_img = np.expand_dims(torch_img, axis=0)

torch_img = torch.Tensor(torch_img.copy()).to(torch.float32)
torch_img = torch_img.to("cuda")

prediction, std_images = predict_single_example(torch_img,
og_ants_img,
config,
models,
overlap,
blend_mode,
tta,
output_std,
fg_bbox)

# Apply postprocessing if required
transforms = ["remove_small_objects", "top_k_cc", "fill_holes"]
for transform in transforms:
if len(config[transform]) > 0:
for i in range(len(config[transform])):
if transform == "remove_small_objects":
transform_kwargs = {"small_object_threshold": config[transform][i][1]}
if transform == "top_k_cc":
transform_kwargs = {"morph_cleanup": config[transform][i][1],
"morph_cleanup_iterations": config[transform][i][2],
"top_k": config[transform][i][3]}
if transform == "fill_holes":
transform_kwargs = {"fill_label": config[transform][i][1]}

prediction = apply_transform(prediction,
transform,
config["labels"],
config[transform][i][0],
transform_kwargs)

# Write prediction mask to nifti file and save to disk
prediction_filename = '{}.nii.gz'.format(str(patient['id']))
output = os.path.join(output_std_dest, prediction_filename)
ants.image_write(prediction, output)

# Write standard deviation image(s) to nifti file and save to disk (only for foreground labels)
if output_std:
for i in range(len(std_images)):
if config["labels"][i] > 0:
std_image_filename = '{}_std_{}.nii.gz'.format(patient['id'], config['labels'][i])
output = os.path.join(output_std_dest, std_image_filename)
ants.image_write(std_images[i], output)

except:
id = patient["id"]
messages += f"Prediction failed for {id}\n"

if len(messages) > 0:
text = Text(messages)
console.print(text)

# Clean up
gc.collect()
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "mist-medical"
version = "0.4.5-alpha"
version = "0.4.6-alpha"
requires-python = ">= 3.8"
description = "MIST is a simple, fully automated framework for 3D medical imaging segmentation."
readme = "README.md"
Expand Down

0 comments on commit 06bc78b

Please sign in to comment.