SIIM-ACR Pneumothorax Segmentation ¶
This is an entry to the Kaggle competition available here.
Initial Analysis of Problem ¶
Pneumothorax: Occurs when air leaks into the space between the lungs and chest wall.
Typically we see the following:
- Black = Air,
- Grey = Fluids & Tissue, and
- White = Bones and other solids.
Indicators in imagery are:
- Air packets,
- Small deflated lung for trauma pneumothorax,
- Blackened lung due to hemopneumothorax,
- Shifted position of heart and trachea away from collapsed lung due to extrinsic pressure for tension pneumothorax.
Air packets will be black in X-ray imagery between the lungs and chest walls. For trauma pneumothorax, it will prevent visibility of vascular markings (blood vessels) going through the lung parenchyma towards the chest wall on the side of the chest where the lung has collapsed. When collapsed lungs are present, we may see asymmetry in the image for one collapsed lung, or we may see two abnormally small lungs if both have collapsed (unlikely as patient would most likely not be able to receive treatment in time to allow for breathing). We may also see smaller (darker) median pixel values for images of patients with pneumothorax.
Noise will occur due to:
- Instrument limitations, and
- Attached medical equipment (e.g. ECG equipment, pacemakers, etc.).
# Set up environment
%matplotlib inline
# Install GitHub repo dependences:
# Pretrained segmentation models with PyTorch.
try:
import segmentation_models_pytorch as smp
except ModuleNotFoundError:
!pip install git+https://github.com/qubvel/segmentation_models.pytorch > /dev/null 2>&1
# Fast image augmentation library.
try:
import albumentations
except ModuleNotFoundError:
!pip install albumentations
# Import packages
import os
import sys
import cv2
import time
import numpy as np
import pandas as pd
import pydicom
from tqdm import tqdm_notebook as tqdm
import matplotlib.pyplot as plt
import matplotlib.patches
import matplotlib.ticker
import seaborn as sns
from albumentations import MultiplicativeNoise, GaussNoise, HorizontalFlip, ShiftScaleRotate, Normalize, Resize, Compose
from albumentations.pytorch import ToTensor
from sklearn.model_selection import StratifiedKFold
# Hide all warnings
import warnings
warnings.filterwarnings('ignore')
# Import DL libraries
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.optim.optimizer import Optimizer, required
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader, Dataset, sampler
import torchvision
# Initialize seeds for reproducibility.
np.random.seed(0)
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Paths
base_dir = '/home/jeffrey/harrison-ai'
train_dir = os.path.join(base_dir, 'train')
test_dir = os.path.join(base_dir, 'test')
sample_submission_path = os.path.join(base_dir, 'stage_2_sample_submission.csv')
train_rle_path = os.path.join(base_dir, 'stage_2_train.csv')
data_folder = os.path.join(base_dir, 'train_png')
test_data_folder = os.path.join(base_dir, 'test_png')
# Import supplied python files from Kaggle competition (stored locally).
sys.path.append(base_dir)
import mask_functions
# Prepare list of complete paths of imagery.
train_img_paths = [os.path.join(train_dir, filename) for filename in os.listdir(train_dir)]
test_img_paths = [os.path.join(test_dir, filename) for filename in os.listdir(test_dir)]
# Load the CSV files with masking information. Note that EncodedPixels = -1 means there is no
# detected pneumothorax in the truth data.
df_train = pd.read_csv(os.path.join(base_dir, 'stage_2_train.csv'))
df_test = pd.read_csv(os.path.join(base_dir, 'stage_2_sample_submission.csv'))
healthy_paths = [os.path.join(train_dir, '{}.dcm'.format(x)) for x in df_train[df_train.EncodedPixels == '-1'].ImageId]
unhealthy_paths = [os.path.join(train_dir, '{}.dcm'.format(x)) for x in df_train[df_train.EncodedPixels != '-1'].ImageId]
# Note that train consists of stage 1 data and test consists of stage 2 data.
# Stage 1 data collected from https://www.kaggle.com/seesee/siim-train-test as it is no
# longer available via the Cloud Healthcare API.
Exploratory Data Analysis ¶
Firstly, let's have a look at the DICOM headers. We can see from below that for our analysis, patient sex and age may be useful variables. Both may relate to the level of noise in the image, e.g. there would typically be more breast tissue for females than males, and older patients may have developed diseases which could cause pneumothoraces. The images are each 1024x1024 (greyscale), which may result in memory issues when training on the dataset. The view position may also be an important feature. After inspection, it is evident that there are two positions: PA and AP CXRs (Chest X-Rays). PA-viewed X-rays pass from posterior to anterior and AP-viewed are those that pass from anterior to posterior. PA CXRs are the standard view. AP CXRs could have been taken instead potentially due to patient immobility. It is typically taken with a patient lying on a bed (i.e. for people who cannot stand upright). This means that the fluid level in the pleural space would be more dispersed across the image for AP CXRs relative to PA CXRs.
# Quick look at the first DICOM metadata.
img = pydicom.dcmread(train_img_paths[0])
print(img) # Print the metadata of the first frame
From the below, we can see that images come with various saturation levels, with some highlighting bone moreso than the tissue. The middle image in the top row shows highly saturated/noisy data.
def plot_img_sample(train_img_paths, ncols=5, nrows=2):
"""Plot the first ncols*nrows images."""
fig, axes = plt.subplots(figsize=(20, 10), ncols=ncols, nrows=nrows, sharex=True, sharey=True)
for i, path in enumerate(train_img_paths[:nrows*ncols]):
img = pydicom.dcmread(path).pixel_array
axes[i % nrows, i // nrows].imshow(img, cmap=plt.cm.bone)
mask = df_train[df_train.ImageId == os.path.splitext(path)[0].split(os.sep)[-1]].iloc[0]['EncodedPixels']
axes[i % nrows, i // nrows].set_xlabel(
'Pneumothorax') if mask != '-1' else axes[i%nrows, i//nrows].set_xlabel('Healthy')
plt.subplots_adjust(wspace=0, hspace=0)
plot_img_sample(train_img_paths, ncols=5, nrows=2)
Functions¶
# Mask decoding functions
def run_length_decode(rle, height=1024, width=1024, fill_value=1):
"""Decodes mask.
From https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
"""
component = np.zeros((height, width), np.float32)
component = component.reshape(-1)
rle = np.array([int(s) for s in rle.strip().split(' ')])
rle = rle.reshape(-1, 2)
start = 0
for index, length in rle:
start = start+index
end = start+length
component[start: end] = fill_value
start = end
component = component.reshape(width, height).T
return component
def run_length_encode(component):
"""Converts mask back to RLE-standard for submission.
From https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
"""
component = component.T.flatten()
start = np.where(component[1:] > component[:-1])[0]+1
end = np.where(component[:-1] > component[1:])[0]+1
length = end-start
rle = []
for i in range(len(length)):
if i == 0:
rle.extend([start[0], length[0]])
else:
rle.extend([start[i]-end[i-1], length[i]])
rle = ' '.join([str(r) for r in rle])
return rle
# Plotting functions
def get_bounding_box(img):
"""Returns the bounding box edges for an image.
"""
rows = np.any(img, axis=1)
cols = np.any(img, axis=0)
ymin, ymax = np.where(rows)[0][[0, -1]]
xmin, xmax = np.where(cols)[0][[0, -1]]
return xmin, ymin, xmax, ymax
def plot_imgs_and_masks(unhealthy_paths, df_train, n_images=5):
fig, axes = plt.subplots(figsize=(20, 10), ncols=n_images, nrows=2, sharex=True, sharey=True)
for i, path in enumerate(unhealthy_paths[:n_images]):
image_name = os.path.splitext(path)[0].split(os.sep)[-1]
img = pydicom.dcmread(path)
# Row 0: Image
axes[0, i].imshow(img.pixel_array, cmap=plt.cm.bone)
# Row 1: Image with mask and bounding box.
mask_encoded = df_train[df_train.ImageId == image_name].iloc[0]['EncodedPixels']
if mask_encoded != '-1':
mask = mask_functions.rle2mask(mask_encoded, 1024, 1024).T # Get mask image
xmin, ymin, xmax, ymax = get_bounding_box(mask) # Get bounding box edges from mask
# Create bounding box.
patch = matplotlib.patches.Rectangle(
(xmin, ymin),
xmax-xmin,
ymax-ymin,
linewidth=1,
edgecolor='r',
facecolor='none'
)
# Plot image with mask and bounding box
axes[1, i].imshow(img.pixel_array, cmap=plt.cm.bone)
axes[1, i].imshow(mask, alpha=0.3, cmap="Reds") # Add mask overlay with 0.3 opacity
axes[1, i].add_patch(patch) # Add bounding box
axes[1, i].set_xlabel(
'Pneumothorax') if mask != '-1' else axes[i%nrows, i//nrows].set_xlabel('Healthy')
plt.subplots_adjust(wspace=0, hspace=0)
Below we observe what pneumothorax looks like by plotting the masks (with some opacity) on the image. It is hard to detect to the untrained eye.
# Initial plot of the first few images with masks for comparison.
plot_imgs_and_masks(unhealthy_paths, df_train, n_images=5)
Safety Precautions¶
Firstly, we must check if there is anything to be cautious of with regards to the data themselves.
It is identified that there are some images without any associated row in the CSV file. Furthermore, there are 12,089 images in the training set and 12,954 associated annotations, meaning there can be multiple masks per image. Therefore, as a safety measure, I checked to see if there were any images which had both a row stating that there is a pneumothorax indicator in the image, and another row stating the same image is free from pneumothorax indicators. It is confirmed that there are no such cases, and the dataset has already been well-processed. Also checked were if any image supplied did not have a row in the CSV file containing the truth data. This occurred for 42 images in the training set. Therefore, there was no reason for the images to be provided.
The below also shows that there are more healthy patients in the training set than those with pneumothorax. Additionally, pneumothoraces do not dominate the images, and so there is a bias in both the images and the dataset as a whole towards the healthy patient. Empirical models develop a bias towards the more dense regions of the distribution in feature space, and as such this bias towards the healthy patients will be a problem that needs to be investigated, however, due to time constraints, this will not be attempted in this report. Ideally, some form of weighting for the pneumothorax indicators (and against that of healthy patients) would be used. Sampling the images to create an equal dataset of images with and without pneumothraces would also be ideal.
num_healthy = len(df_train.EncodedPixels[df_train.EncodedPixels == '-1'])
num_unhealthy = len(df_train.EncodedPixels[df_train.EncodedPixels != '-1'])
print(
'Total number of images in training set: {}\n'
'Total number of annotations in training set: {}\n'
'Number of healthy patients: {}\n'
'Number of patients with pneumothorax: {}\n\n'
'Total number of images in test set: {}\n'
'Total number of annotations in test set: {}'
''.format(
len(train_img_paths),
len(df_train),
num_healthy,
num_unhealthy,
len(test_img_paths),
len(df_test))
)
def check_for_mislabelled_data(df, dataset_name=''):
"""Checks if any images with multiple masks have a value of '-1'
which indicates that no mask is present."""
df_many_masks = df.groupby('ImageId').count()
df_many_masks = df_many_masks[df_many_masks.EncodedPixels > 1]
count = 0 # Count of number of images with errors in labelling.
for img_id in df_many_masks.index.values:
for idx, row in df[df.ImageId == img_id].iterrows():
if row.EncodedPixels == '-1':
print(df[df.ImageId == img_id])
count += 0
print('{} images have been mislabelled from the {} dataset.'.format(count, dataset_name))
def check_for_missing_annotations(path, df):
"""Checks if there are missing annotations for any of the imagery in path."""
missing_annotations = []
for img_path in path:
img_filename = os.path.splitext(img_path)[0].split(os.sep)[-1]
if not df_train.ImageId.str.contains(img_filename).any():
missing_annotations.append(img_filename)
return missing_annotations
check_for_mislabelled_data(df_train, dataset_name='train')
train_missing_ann = check_for_missing_annotations(df=df_train, path=train_img_paths)
print('Number of images with missing annotations:\n'
'Train: {}'.format(len(train_missing_ann)))
Next, the maximum number of annotations is inspected. 10 annotations occurs in only 1 image, with the plot below showing that the majority have only one or two masks for images with pneumothoraces.The multiple annotations in a single image suggest either multiple instances of pneumothoraces indicators, or multiple segmented indicators of the condition are present in the image, possibly separated by noise. As we are looking at several different types of pneumothoraces, the indicators of the more severe and rarer types (e.g. hemopneumothoraces and trauma pneumothoraces) might not be truly represented by the trained model. For instance, with hemopneumothoraces, the fluid level is a straight line, but other fluid there is a meniscus so it curls upward.
def image_diagnostics(df_train, n=10):
"""Plots the distribution of number of annotations in each image and prints
the maximum number of annotations made on a single image and n images with
the largest number of annotations."""
df_train_grouped = df_train.groupby('ImageId').count()
df_train_grouped = df_train_grouped.sort_values(by='EncodedPixels', ascending=False)
print(
'Maximum number of annotations on image:'
'{} \n\nTop {} images with the most annotations:\n'
'{}'.format(
df_train_grouped.EncodedPixels.max(),
n,
df_train_grouped.EncodedPixels.head(n).to_string())
)
fig, ax = plt.subplots()
sns.distplot(df_train_grouped, kde=False)
loc = matplotlib.ticker.MultipleLocator(base=1.0)
ax.xaxis.set_major_locator(loc)
ax.set_xlabel('Number of Annotations in Image')
ax.set_ylabel('Frequency')
ax.set_title('Distribution of Annotations in Images')
plt.show()
image_diagnostics(df_train, n=10)
# Add some additional information from the images for further exploratory data analysis.
# Note that sex and age are considered for each image rather than each person as this
# will be relative information to the bias in patient selection.
def add_df_info(df, directory):
"""Adds the age, sex, health (boolean), mask area size and view position
information to the dataframe:"""
healthy = []
sex = []
age = []
mask_size = []
view = []
for idx, row in tqdm(df.iterrows(), desc='Adding sex, age and health to df'):
if row.EncodedPixels in (-1, '-1'):
healthy.append(True)
mask_size.append(0)
else:
healthy.append(False)
mask_size.append(np.count_nonzero(run_length_decode(row.EncodedPixels)))
# Read in image to get sex and age data.
img = pydicom.dcmread(os.path.join(directory, '{}.dcm'.format(row.ImageId)))
sex.append(img.data_element('PatientSex').value)
age.append(img.data_element('PatientAge').value)
view.append(img.data_element('ViewPosition').value)
df = df.assign(
Sex=sex,
Age=age,
Healthy=healthy,
MaskSize=mask_size,
View=view
)
df.Sex = df.Sex.astype('category')
df.Age = df.Age.astype('int')
df.MaskSize = df.MaskSize.astype('int')
return df
df_train = add_df_info(df=df_train, directory=train_dir)
df_test = add_df_info(df=df_test, directory=test_dir)
Load/Save Data¶
# # Save dfs for faster loading.
# df_train.to_csv(os.path.join(base_dir, 'df_train.csv'), index=False)
# df_test.to_csv(os.path.join(base_dir, 'df_test.csv'), index=False)
# # Read dfs for faster loading.
df_train = pd.read_csv(os.path.join(base_dir, 'df_train.csv'))
df_test = pd.read_csv(os.path.join(base_dir, 'df_test.csv'))
Further Analysis of Metadata¶
From the top figure below, we can see that there are somewhat more men than women in the dataset for both healthy patients and those with pneumothoraces. Both males and females have less cases of images with pneumothoraces in comparison to those without it.
Furthermore, there are two instances where the patient's age is above 100 (age=148, age=413). These are both clearly outliers which should be removed if the age feature is used for analysis.
Overall, the age distribution is somewhat Gaussian for healthy female patients and those with pneumothorax indicators. There appears to be a peak between 50-60 years of age and slight left skew, which may both be due to a natural effect due to pneumothorax being correlated with lung disease and mechanical ventilation, both of which are more likely as age increases.
Interestingly, there is a bimodal distribution for male patients with pneumothoraces, with the peaks at 20-30 and 50-60 years of age. This suggests that young male patients (<30 years old) may have some correlationship with developing pneumothoraces. Withers et al. (1964) previously identified this relationship with spontaneous pneumothoraces. This suggests that this might be the type of pneumothoraces predominantly found in the first peak region of the age distribution.
print('List of patients with age greater than or equal to 100 or less than 0:\n'
'{}'.format(df_train[(df_train.Age >= 100) | df_train.Age < 0)][['ImageId', 'Sex', 'Age', 'Healthy']]))
def plot_metadata_info(df_train):
"""Plots the age distribution for both sexes, as well as the distribution of both sexes, for patients with
and without a pneumothorax indicator present in their X-ray."""
fig = plt.figure(figsize=(10, 10))
gs = fig.add_gridspec(2, 2)
ax1 = fig.add_subplot(gs[0, :])
sns.countplot(
x='Sex',
hue='Healthy',
data=df_train[['Age', 'Sex', 'Healthy']],
palette='rainbow',
order=['M', 'F']
)
ax1.set(title='Number of Patients with and without Pneumothorax for both Sexes')
ax2 = fig.add_subplot(gs[1, 0])
sns.distplot(
df_train[(df_train.Healthy) & (df_train.Age < 100) & (df_train.Sex == 'M')]['Age'],
label='Healthy'
)
sns.distplot(
df_train[(~df_train.Healthy) & (df_train.Age < 100) & (df_train.Sex == 'M')]['Age'],
label='Pneumothorax'
)
ax2.set(xlabel='Age', ylabel='Normalised Count', title='Male Age Distribution')
ax2.set_xlim(xmin=0, xmax=100)
ax2.legend()
ax3 = fig.add_subplot(gs[1, 1])
sns.distplot(
df_train[(df_train.Healthy) & (df_train.Age < 100) & (df_train.Sex == 'F')]['Age'],
label='Healthy'
)
sns.distplot(
df_train[(~df_train.Healthy) & (df_train.Age < 100) & (df_train.Sex == 'F')]['Age'],
label='Pneumothorax'
)
ax3.set(xlabel='Age', title='Female Age Distribution')
ax3.set_xlim(xmin=0, xmax=100)
ax3.legend()
plt.show()
plot_metadata_info(df_train)
The area of each mask is also a useful piece of analysis. It can be seen in the figure below that the mask area distribution is Poisson. As the model developed will not be perfect, some individual pixels, or small groups of pixels (which are influenced moreso by noise than data) will be mislabelled as pneumothorax indicators. Therefore, a minimum area threshold will be required to filter out some small false positives due to noise and therefore develop the optimal predictions for the test set.
def plot_area_histogram(df_train, binwidth=1000):
plt.figure(figsize=(15, 5))
x = df_train[~df_train.Healthy].MaskSize
plt.hist(x, bins=np.arange(0, np.max(x) + binwidth, binwidth))
plt.title('Frequency of Mask Area of Pneumothorax (bins=1000 pix)')
plt.xlabel('Area')
plt.ylabel('Frequency')
plt.xlim(xmin=0)
plt.gca().xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(10000))
plot_area_histogram(df_train, binwidth=1000)
A heatmap is generated of where the pneumothorax masked pixels are found in the training set.
With both views, the top of the lungs is where the masked pixels are more commonly found. This could be due to it being an easier region to identify the pneumothorax indicators, either due to less noise from breast tissues and other organs, or it could be due to other reasons which would need to be identified through discussions with experts.
This is particularly evident with the PA-viewed X-rays, where pneumothoraces are not identified in the middle-region of the centre of the chest. It is evident that the AP-viewed images are more well distributed and clearly show the shapes of the lungs. Therefore, this suggests that experts can better detect and annotate pneumothoraces in the AP-viewed images rather than the PA-viewed images.
The left-hand side appears to be brighter for images taken from both view, which is the left lun (as the mask was flipped). This may suggest that the additional pressure from the heart on that side of the chest could potentially cause more collapsed lungs.
# Create heat map of location of pneumothorax indicators.
def create_heatmaps(df_train):
"""Generates heatmaps for both AP and PA views."""
ap_heatmap = np.zeros((1024, 1024), dtype=np.float64)
pa_heatmap = np.zeros((1024, 1024), dtype=np.float64)
for idx, row in tqdm(df_train.iterrows()):
if row.EncodedPixels not in (-1, '-1'):
mask = mask_functions.rle2mask(row.EncodedPixels, 1024, 1024)
mask = np.rot90(mask, 3) #rotating three times 90 to the right place
mask = np.flip(mask, axis=1)
if row.View == 'AP':
ap_heatmap += mask
elif row.View == 'PA':
pa_heatmap += mask
else:
print('{} is not either the AP or PA view.'.format(row.View))
# Normalize ap_heatmap and pa_heatmap
max_pixel_value = max(np.max(pa_heatmap), np.max(ap_heatmap))
fig, axes = plt.subplots(figsize=(15, 5), nrows=1, ncols=2)
cm_0 = axes[0].imshow(ap_heatmap, vmin=0, vmax=max_pixel_value, cmap='hot', interpolation='nearest')
axes[0].set_title('AP-viewed Image Heatmap')
cm_1 = axes[1].imshow(pa_heatmap, vmin=0, vmax=max_pixel_value, cmap='hot', interpolation='nearest')
axes[1].set_title('PA-viewed Image Heatmap')
fig.colorbar(cm_1, ax=axes.ravel().tolist())
plt.show()
create_heatmaps(df_train)
Model Selection¶
Existing Approaches¶
For the model we observe that of the existing notebooks, mobassir achieved a public score of 0.9006 with a feature pyramid network (FPN) with InceptionResNetv2 encoder (using 512x512 images). meaninglesslives obtained a public score of 0.8997 with UNet++ with EfficientNet-B4 encoder pretrained on ImageNet (using 256x256 images). It would be interesting to see what improvements could be done by using EfficientNet-B7 as the backend, and by using 512x512 images. This network is deeper than the B4 model and results in (marginal) improvements on model accuracy. However, the memory required renders it unable to be tested with my current working environment (it would require approval for GPU clusters, which is currently pending). This will probably be finished before showcasing my results, in which case I plan on testing it afterwards regardless.
As such, two different models are tested. An FPN with EfficientNet-B0 encoder is used with 1024x1024 images and batch size of 2, and an FPN with EfficientNet-B3 encoder is used with 512x512 images and batch size of 4.
Prepare dataset class¶
class SIIMDataset(Dataset):
"""PyTorch-appropriate dataset class for the
SIIM Pneumothorax Kaggle challenge.
Adaptations made from intial code at:
https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
"""
def __init__(self, df, fnames, data_folder, size, mean, std, phase):
self.df = df
self.root = data_folder
self.size = size
self.mean = mean
self.std = std
self.phase = phase
self.transforms = get_transforms(phase, size, mean, std)
self.gb = self.df.groupby('ImageId')
self.fnames = fnames
def __getitem__(self, idx):
image_id = self.fnames[idx]
df = self.gb.get_group(image_id)
annotations = df['EncodedPixels'].tolist()
image_path = os.path.join(self.root, image_id + ".png")
image = cv2.imread(image_path)
mask = np.zeros([1024, 1024])
if annotations[0] != '-1':
for rle in annotations:
mask += run_length_decode(rle)
mask = (mask >= 1).astype('float32')
augmented = self.transforms(image=image, mask=mask)
image = augmented['image']
mask = augmented['mask']
return image, mask
def __len__(self):
return len(self.fnames)
def get_transforms(phase, size, mean, std):
"""From https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2"""
list_transforms = []
if phase == "train":
list_transforms.extend(
[
HorizontalFlip(p=0.5),
ShiftScaleRotate(
shift_limit=0, # No resizing
scale_limit=0.1,
rotate_limit=10, # Rotate
p=0.5,
border_mode=cv2.BORDER_CONSTANT
),
GaussNoise(),
MultiplicativeNoise(multiplier=1.5, p=1),
]
)
list_transforms.extend(
[
Resize(size, size),
Normalize(mean=mean, std=std, p=1),
ToTensor(),
]
)
list_trfms = Compose(list_transforms)
return list_trfms
def provider(fold, total_folds, data_folder, df_path, phase, size,
mean=None, std=None, batch_size=8, num_workers=2):
"""Adaptations made from intial code at:
https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
"""
df_all = pd.read_csv(df_path)
df = df_all.drop_duplicates('ImageId')
df_with_mask = df[df["EncodedPixels"] != "-1"]
df_with_mask['has_mask'] = 1
df_without_mask = df[df["EncodedPixels"] == "-1"]
df_without_mask['has_mask'] = 0
df_without_mask_sampled = df_without_mask.sample(len(df_with_mask)+1500, random_state=0) # random state is imp
df = pd.concat([df_with_mask, df_without_mask_sampled])
# Note that an equal number of positive and negative cases are chosen.
kfold = StratifiedKFold(total_folds, shuffle=True, random_state=0)
train_idx, val_idx = list(kfold.split(df["ImageId"], df["has_mask"]))[fold]
train_df, val_df = df.iloc[train_idx], df.iloc[val_idx]
df = train_df if phase == "train" else val_df
fnames = df['ImageId'].values
image_dataset = SIIMDataset(df_all, fnames, data_folder, size, mean, std, phase)
dataloader = DataLoader(
image_dataset,
batch_size=batch_size,
num_workers=num_workers,
pin_memory=True,
shuffle=True,
)
return dataloader
Defining the Loss Function¶
The notebook that this project is based off uses a combination of dice loss and focal loss. Dice loss is the evaluation metric used in the competition and focal loss performs well in class-imbalanced problems such as the problem at hand. For consistency, we use the same loss function.
# Loss functions from
# https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
def dice_loss(input, target):
input = torch.sigmoid(input)
smooth = 1.0
iflat = input.view(-1)
tflat = target.view(-1)
intersection = (iflat * tflat).sum()
return ((2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))
class FocalLoss(nn.Module):
def __init__(self, gamma):
super().__init__()
self.gamma = gamma
def forward(self, input, target):
if not (target.size() == input.size()):
raise ValueError("Target size ({}) must be the same as input size ({})"
.format(target.size(), input.size()))
max_val = (-input).clamp(min=0)
loss = input - input * target + max_val + \
((-max_val).exp() + (-input - max_val).exp()).log()
invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))
loss = (invprobs * self.gamma).exp() * loss
return loss.mean()
class MixedLoss(nn.Module):
def __init__(self, alpha, gamma):
super().__init__()
self.alpha = alpha
self.focal = FocalLoss(gamma)
def forward(self, input, target):
loss = self.alpha*self.focal(input, target) - torch.log(dice_loss(input, target))
return loss.mean()
# RAadam optimizer from
# https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
class RAdam(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
self.buffer = [[None, None, None] for ind in range(10)]
super(RAdam, self).__init__(params, defaults)
def __setstate__(self, state):
super(RAdam, self).__setstate__(state)
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data.float()
if grad.is_sparse:
raise RuntimeError('RAdam does not support sparse gradients')
p_data_fp32 = p.data.float()
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p_data_fp32)
state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
else:
state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
exp_avg.mul_(beta1).add_(1 - beta1, grad)
state['step'] += 1
buffered = self.buffer[int(state['step'] % 10)]
if state['step'] == buffered[0]:
N_sma, step_size = buffered[1], buffered[2]
else:
buffered[0] = state['step']
beta2_t = beta2 ** state['step']
N_sma_max = 2 / (1 - beta2) - 1
N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
buffered[1] = N_sma
# more conservative since it's an approximated value
if N_sma >= 5:
step_size = group['lr'] * np.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])
else:
step_size = group['lr'] / (1 - beta1 ** state['step'])
buffered[2] = step_size
if group['weight_decay'] != 0:
p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
# more conservative since it's an approximated value
if N_sma >= 5:
denom = exp_avg_sq.sqrt().add_(group['eps'])
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
else:
p_data_fp32.add_(-step_size, exp_avg)
p.data.copy_(p_data_fp32)
return loss
# Utility functions from
# https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2
def predict(X, threshold):
X_p = np.copy(X)
preds = (X_p > threshold).astype('uint8')
return preds
def metric(probability, truth, threshold=0.5, reduction='none'):
'''Calculates dice of positive and negative images seperately'''
'''probability and truth must be torch tensors'''
batch_size = len(truth)
with torch.no_grad():
probability = probability.view(batch_size, -1)
truth = truth.view(batch_size, -1)
assert(probability.shape == truth.shape)
p = (probability > threshold).float()
t = (truth > 0.5).float()
t_sum = t.sum(-1)
p_sum = p.sum(-1)
neg_index = torch.nonzero(t_sum == 0)
pos_index = torch.nonzero(t_sum >= 1)
dice_neg = (p_sum == 0).float()
dice_pos = 2 * (p*t).sum(-1)/((p+t).sum(-1))
dice_neg = dice_neg[neg_index]
dice_pos = dice_pos[pos_index]
dice = torch.cat([dice_pos, dice_neg])
dice_neg = np.nan_to_num(dice_neg.mean().item(), 0)
dice_pos = np.nan_to_num(dice_pos.mean().item(), 0)
dice = dice.mean().item()
num_neg = len(neg_index)
num_pos = len(pos_index)
return dice, dice_neg, dice_pos, num_neg, num_pos
class Meter:
'''A meter to keep track of iou and dice scores throughout an epoch'''
def __init__(self, phase, epoch):
self.base_threshold = 0.5
self.base_dice_scores = []
self.dice_neg_scores = []
self.dice_pos_scores = []
self.iou_scores = []
def update(self, targets, outputs):
probs = torch.sigmoid(outputs)
dice, dice_neg, dice_pos, _, _ = metric(probs, targets, self.base_threshold)
self.base_dice_scores.append(dice)
self.dice_pos_scores.append(dice_pos)
self.dice_neg_scores.append(dice_neg)
preds = predict(probs, self.base_threshold)
iou = compute_iou_batch(preds, targets, classes=[1])
self.iou_scores.append(iou)
def get_metrics(self):
dice = np.mean(self.base_dice_scores)
dice_neg = np.mean(self.dice_neg_scores)
dice_pos = np.mean(self.dice_pos_scores)
dices = [dice, dice_neg, dice_pos]
iou = np.nanmean(self.iou_scores)
return dices, iou
def epoch_log(phase, epoch, epoch_loss, meter, start):
'''logging the metrics at the end of an epoch'''
dices, iou = meter.get_metrics()
dice, dice_neg, dice_pos = dices
print("Loss: %0.4f | dice: %0.4f | dice_neg: %0.4f | dice_pos: %0.4f | IoU: %0.4f" % (epoch_loss, dice, dice_neg, dice_pos, iou))
return dice, iou
def compute_ious(pred, label, classes, ignore_index=255, only_present=True):
'''computes iou for one ground truth mask and predicted mask'''
pred[label == ignore_index] = 0
ious = []
for c in classes:
label_c = label == c
if only_present and np.sum(label_c) == 0:
ious.append(np.nan)
continue
pred_c = pred == c
intersection = np.logical_and(pred_c, label_c).sum()
union = np.logical_or(pred_c, label_c).sum()
if union != 0:
ious.append(intersection / union)
return ious if ious else [1]
def compute_iou_batch(outputs, labels, classes=None):
'''computes mean iou for a batch of ground truth masks and predicted masks'''
ious = []
preds = np.copy(outputs) # copy is imp
labels = np.array(labels) # tensor to np
for pred, label in zip(preds, labels):
ious.append(np.nanmean(compute_ious(pred, label, classes)))
iou = np.nanmean(ious)
return iou
class Trainer(object):
'''Train and validate the model.
Adapted from:
https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2.
'''
def __init__(self, model, model_name='unet', size=512, batch_size=4, num_workers=4):
self.model_name = model_name
self.fold = 1
self.total_folds = 5
self.num_workers = num_workers
self.batch_size = {"train": batch_size, "val": batch_size}
self.accumulation_steps = 32 // self.batch_size['train']
self.lr = 5e-4
self.num_epochs = 32
self.best_loss = float("inf")
self.phases = ["train", "val"]
self.device = torch.device("cuda:0")
torch.set_default_tensor_type("torch.cuda.FloatTensor")
self.net = model
self.criterion = MixedLoss(10.0, 2.0)
#self.optimizer = optim.Adam(self.net.parameters(), lr=self.lr)
self.optimizer = RAdam(model.parameters(), lr=self.lr)
self.scheduler = ReduceLROnPlateau(self.optimizer, mode="min", patience=3, verbose=True)
self.net = self.net.to(self.device)
cudnn.benchmark = True
self.dataloaders = {
phase: provider(
fold=1,
total_folds=5,
data_folder=data_folder,
df_path=train_rle_path,
phase=phase,
size=size,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
batch_size=self.batch_size[phase],
num_workers=self.num_workers,
)
for phase in self.phases
}
self.losses = {phase: [] for phase in self.phases}
self.iou_scores = {phase: [] for phase in self.phases}
self.dice_scores = {phase: [] for phase in self.phases}
def forward(self, images, targets):
images = images.to(self.device)
masks = targets.to(self.device)
outputs = self.net(images)
loss = self.criterion(outputs, masks)
return loss, outputs
def iterate(self, epoch, phase):
meter = Meter(phase, epoch)
start = time.strftime("%H:%M:%S")
print(f"Starting Epoch: {epoch} | Phase: {phase} | Current Time: {start}")
batch_size = self.batch_size[phase]
self.net.train(phase == "train")
dataloader = self.dataloaders[phase]
running_loss = 0.0
total_batches = len(dataloader)
self.optimizer.zero_grad()
for itr, batch in enumerate(dataloader):
images, targets = batch
loss, outputs = self.forward(images, targets)
loss = loss / self.accumulation_steps
if phase == "train":
loss.backward()
if (itr + 1 ) % self.accumulation_steps == 0:
self.optimizer.step()
self.optimizer.zero_grad()
running_loss += loss.item()
outputs = outputs.detach().cpu()
meter.update(targets, outputs)
epoch_loss = (running_loss * self.accumulation_steps) / total_batches
dice, iou = epoch_log(phase, epoch, epoch_loss, meter, start)
self.losses[phase].append(epoch_loss)
self.dice_scores[phase].append(dice)
self.iou_scores[phase].append(iou)
torch.cuda.empty_cache()
return epoch_loss
def start(self):
for epoch in tqdm(range(self.num_epochs)):
self.iterate(epoch, "train")
state = {
"epoch": epoch,
"best_loss": self.best_loss,
"state_dict": self.net.state_dict(),
"optimizer": self.optimizer.state_dict(),
}
val_loss = self.iterate(epoch, "val")
self.scheduler.step(val_loss)
if val_loss < self.best_loss:
print("******** New optimal found, saving state ********")
state["best_loss"] = self.best_loss = val_loss
torch.save(state, "./{}_model.pth".format(self.model_name))
print()
def plot_training_metrics(model_trainer):
"""Plots the losses, dice scores and IoU scores of the model
throughout the training process."""
fig, axes = plt.subplots(figsize=(15, 15), ncols=1, nrows=3)
names = ('Losses', 'Dice Scores', 'IOU Scores')
for i, metric in enumerate([
model_trainer.losses,
model_trainer.dice_scores,
model_trainer.iou_scores
]):
axes[i].plot(
range(len(metric["train"])),
metric["train"],
label=f'train {names[i]}'
)
axes[i].plot(
range(len(metric["train"])),
metric["val"],
label=f'val {names[i]}'
)
axes[i].set_title(f'{names[i]} plot')
plt.xlabel('Epoch')
plt.ylabel(f'{names[i]}')
axes[i].set_xlabel('Epoch')
axes[i].set_ylabel('Dice Loss')
axes[i].legend(loc='upper left')
class TestDataset(Dataset):
""" Test dataset.
Adapted from:
https://www.kaggle.com/mobassir/unet-with-se-resnext50-32x4d-encoder-for-stage-2.
'''"""
def __init__(self, root, df, size, mean, std, tta=4):
self.root = root
self.size = size
self.fnames = list(df["ImageId"])
self.num_samples = len(self.fnames)
self.transform = Compose(
[
Normalize(mean=mean, std=std, p=1),
Resize(size, size),
ToTensor(),
]
)
def __getitem__(self, idx):
fname = self.fnames[idx]
path = os.path.join(self.root, fname + ".png")
image = cv2.imread(path)
images = self.transform(image=image)["image"]
return images
def __len__(self):
return self.num_samples
def post_process(probability, threshold, min_size):
mask = cv2.threshold(
probability, threshold, 1, cv2.THRESH_BINARY)[1]
num_component, component = cv2.connectedComponents(
mask.astype(np.uint8))
predictions = np.zeros((1024, 1024), np.float32)
num = 0
for c in range(1, num_component):
p = (component == c)
if p.sum() > min_size:
predictions[p] = 1
num += 1
return predictions, num
def save_submission(model_trainer, model_name,
size=512, mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
num_workers=8, batch_size=16,
best_threshold=0.5, min_size=3500):
device = torch.device("cuda:0")
df = pd.read_csv(sample_submission_path)
testset = DataLoader(
TestDataset(test_data_folder, df, size, mean, std),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
)
model = model_trainer.net # get the model from model_trainer object
model.eval()
state = torch.load('./{}_model.pth'.format(model_name), map_location=lambda storage, loc: storage)
model.load_state_dict(state["state_dict"])
encoded_pixels = []
for i, batch in enumerate(tqdm(testset)):
preds = torch.sigmoid(model(batch.to(device)))
preds = preds.detach().cpu().numpy()[:, 0, :, :] # (batch_size, 1, size, size) -> (batch_size, size, size)
for probability in preds:
if probability.shape != (1024, 1024):
probability = cv2.resize(probability, dsize=(1024, 1024), interpolation=cv2.INTER_LINEAR)
predict, num_predict = post_process(probability, best_threshold, min_size)
if num_predict == 0:
encoded_pixels.append('-1')
else:
r = run_length_encode(predict)
encoded_pixels.append(r)
df['EncodedPixels'] = encoded_pixels
df.to_csv('{}_submission.csv'.format(model_name), columns=['ImageId', 'EncodedPixels'], index=False)
return df
Train Models ¶
FPN Model with EfficientNet-B0 Encoder and 1024x1024 Images (batch size = 2)¶
fpn_efficientb0_model = smp.FPN('efficientnet-b0', encoder_weights="imagenet")
print(fpn_efficientb0_model)
fpn_efficientb0_model_trainer = Trainer(fpn_efficientb0_model, model_name='fpn_efficientb0', size=1024, batch_size=2, num_workers=16)
fpn_efficientb0_model_trainer.start()
# Plot training performance for FPN model
plot_training_metrics(fpn_efficientb0_model_trainer)
# Save FPN model results
df_fpn = save_submission(fpn_efficientb0_model_trainer, model_name='fpn_efficientb0',
num_workers=8, batch_size=4, best_threshold=0.5, min_size=3500)
FPN Model with EfficientNet-B3 Encoder with 512x512 Images (batch size = 4)¶
fpn_efficientb3_model = smp.FPN('efficientnet-b3', encoder_weights="imagenet")
print(fpn_efficientb3_model)
fpn_efficientb3_model_trainer = Trainer(fpn_efficientb3_model, model_name='fpn_efficient', size=512, batch_size=4, num_workers=4)
fpn_efficientb3_model_trainer.start()
# Plot training performance for FPN model
plot_training_metrics(fpn_efficientb3_model_trainer)
# Save FPN model results
df_fpn_efficientb3 = save_submission(fpn_efficientb3_model_trainer, model_name='fpn_efficientb3',
num_workers=8, batch_size=4, best_threshold=0.5, min_size=3500)
Final Remarks¶
Shortcomings¶
Due to the time constraints, I have been unable to develop better feature preprocessing. Hardware limitations also affected model performance. A single NVIDIA 1080 TI was used throughout the pipeline (11 GB memory). The batch size was preferred to be at least 4 images per batch due to each batch representing a sample population. The skewed dataset towards healthy patients means that a significant portion of batches would have no masked pixels, which is problematic in minimising bias. Note that each model discussed used an FPN with EfficientNet encoder.
Downsampling the images from 1024x1024 to 512x512 allowed for using a deeper EfficientNet (EfficientNet-B3) due to memory limitations. Using the deeper encoder (EfficientNet-B3) but with downsampled images of 512x512 and batch size of 4 rather than 2 resulted in a private score of 0.8280. This translated to a public score of 0.8998 on the test set once submitted to the Kaggle competition. The FPN with EfficientNet-B0 without downsampled images (i.e. 1024x1024 dimensions) resulted in a private score of 0.7910, which translated to a public score of 0.8411. This shows that the effect of using the shallower backend (EfficientNet-B0 rather than EfficientNet-B3) in combination with batch size of 4 images affects the model performance moreso than the combination of downsampling from 1024x1024 images to 512x512 and using a batch size of 2 images.
Improvements to the Dataset¶
The dataset could be improved by labelling the different types of pneumothoraces found in the images. This could also help speed up the time for a complete diagnosis, which is particuarly important for the less obvious but still fatal types of pneumothoraces such as hemopneumothoraces. It could also be used to identify if young males are truly at risk of tension pneumothoraces in particular or if it is another type. Height also appears to be an indicator in existing literature for this type of pneumothorax and so could also be useful. Including multiple pneumothorax classes would also enable better post-model analysis to determine which types of pneumothoraces the model fail to detect to a suitable level. For instance, is the model able to capture the miniscal lining for lungs filled with air, or the straight lines of the fluid level due to hemopneumothoraces. Whilst entries in the Kaggle competition have yielded results of even 93% accuracy, the data labelling limitations could be preventative of more accurate models.
Additional data would of course be useful, particularly to level out the class discrepency with regards to the images as a whole being either healthy or containing a pneumothorax. The training set was limited to offset this discrepency, which is not ideal at all.
With more time I would have also liked to look at incorporating metadata and some forms of preprocessing to help improve the results, i.e. by trying to manually detect the likelihood of mediastinal shift by comparing both halves of the images together. In more detail:
- Use a Hough line transformation to identify the spinal cord.
- Horizontally flip the right hand side of the image relative to the spinal cord, overlay over the left hand side of the image and subtract the pixels from each other.
- This assumes that there are sharp changes where pneumothorax indicators are present and could particularly be helpful for tension pneumothorax (e.g. due to dark pixels (air) where a collapsed lung is present).
- Copy, horizontally flip and adjoin the subtracted image to itself so that it represents the complete body oncemore rather than just half. Pad images with 0s to create identically shaped imagery.
- Note that this would need to accommodate for noise caused by medical devices, etc.
Furthermore, I'd like to use the median pixel intensity to help detect if there is a significant build up of air in the body as these pixels will be darker, and as such the median pixel intensity would be lower for patients with pneumothorax. This indicator could have issues due to different contrasts used across the imagery. Identifying background vs foreground and their ratios in terms of number of pixels and pixel intensity could help make this a useful indicator.
Furthermore, we must note that the data was supplied by the Society for Imaging Informatics in Medicine, and has been well-privatised. We do not know if the images selected were using advanced X-ray technology, nor if imperfect images were removed. Testing the model on an independent dataset such as that from local hospitals would be needed to determine its usefulness when applied to potentially noisier imagery.
Lastly, the ratio of the height of both diaphram domes at the bottom of each image could be used to indicate external pressure caused by pneumothorax. Typically, the right dome is higher by up to 2 cm as the left is pushed down due to pressure from the heart. If it is below the left, then that could be a useful indicator to include alongside the mask during training.
Improvements to the Evaluation Metric¶
No clear improvements can be made with regards to the evaluation metric. As the competition is evaluated on the mean Dice coefficient, it is clearly useful for optimising the model to best perform in the competition (rather than in real world application). As the dataset is highly imbalanced against pneumothoraces, the focal loss reduces this issue. Therefore, the mixed loss approach (focal + dice losses) used by the leading notebooks is logically sound.
Improvements to the Algorithm¶
The UNet++ model with EfficientNet-B4 encoder used by meaninglesslives performed extremely well with 256x256 imagery (0.8997 in comparison to the 0.9006 by mobassir with 512x512 images). This model is specifically designed for biomedical imaging (see Zhou et al., 2018: A Nested U-Net Architecture for Medical Image Segmentation for more details). This model is in essense a nested (and thus deeper) version of the original UNet model. I tried both models with a batch size of 2 and with the shallowest EfficientNet backend (B0) and continued to run into memory issues.
As the UNet++ model performed so well with 256x256 images, and testing showed that there is a 5-10% improvement when using 512x512 images instead of 256x256, it would be interesting to see the performance of the model used by meaninglesslives but with 512x512 images or even 1024x1024 and with deeper EfficientNet encoders. Furthermore, the model used by meaninglesslives shows an IoU of between 70-80%, which is significantly more than the 30-40% with the FPN model. This shows that its high performance isn't due to the bias caused by the overwhelmingly high number of pixels which are classified as being healthy, but due to properly encapturing the underlying features of pneumothoraces.
Lastly, optimizing the threshold value for the test set is a simple way to improve model performance results. Here, a simple threshold of 0.5 is used, which could explain the poor performance in terms of IoU for both FPN models. Also, a minimum mask area of 3500 is used (as was used by the notebook used). Time limitations have prevented both of these two optimizations from being considered and would be the first on the list of steps required to improve on the model performance.