"""
visualize.py
=======================
Plots SHAP outputs, UMAP embeddings, and overlays predictions on top of WSI.
"""
import plotly.graph_objs as go
import plotly.offline as py
import pandas as pd, numpy as np
import networkx as nx
import dask.array as da
from PIL import Image
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
import sqlite3
import seaborn as sns
from os.path import join
from pathflowai.utils import npy2da
sns.set()
class PlotlyPlot:
"""Creates plotly html plots."""
def __init__(self):
self.plots=[]
def add_plot(self, t_data_df, G=None, color_col='color', name_col='name', xyz_cols=['x','y','z'], size=2, opacity=1.0, custom_colors=[]):
"""Adds plotting data to be plotted.
Parameters
----------
t_data_df:dataframe
3-D transformed dataframe.
G:nx.Graph
Networkx graph.
color_col:str
Column to use to color points.
name_col:str
Column to use to name points.
xyz_cols:list
3 columns that denote x,y,z coords.
size:int
Marker size.
opacity:float
Marker opacity.
custom_colors:list
Custom colors to supply.
"""
plots = []
x,y,z=tuple(xyz_cols)
if t_data_df[color_col].dtype == np.float64:
plots.append(
go.Scatter3d(x=t_data_df[x], y=t_data_df[y],
z=t_data_df[z],
name='', mode='markers',
marker=dict(color=t_data_df[color_col], size=size, opacity=opacity, colorscale='Viridis',
colorbar=dict(title='Colorbar')), text=t_data_df[color_col] if name_col not in list(t_data_df) else t_data_df[name_col]))
else:
colors = t_data_df[color_col].unique()
c = sns.color_palette('hls', len(colors))
c = np.array(['rgb({})'.format(','.join(((np.array(c_i)*255).astype(int).astype(str).tolist()))) for c_i in c])#c = ['hsl(' + str(h) + ',50%' + ',50%)' for h in np.linspace(0, 360, len(colors) + 2)]
if custom_colors:
c = custom_colors
color_dict = {name: c[i] for i,name in enumerate(sorted(colors))}
for name,col in color_dict.items():
plots.append(
go.Scatter3d(x=t_data_df[x][t_data_df[color_col]==name], y=t_data_df[y][t_data_df[color_col]==name],
z=t_data_df[z][t_data_df[color_col]==name],
name=str(name), mode='markers',
marker=dict(color=col, size=size, opacity=opacity), text=t_data_df.index[t_data_df[color_col]==name] if 'name' not in list(t_data_df) else t_data_df[name_col][t_data_df[color_col]==name]))
if G is not None:
#pos = nx.spring_layout(G,dim=3,iterations=0,pos={i: tuple(t_data.loc[i,['x','y','z']]) for i in range(len(t_data))})
Xed, Yed, Zed = [], [], []
for edge in G.edges():
if edge[0] in t_data_df.index.values and edge[1] in t_data_df.index.values:
Xed += [t_data_df.loc[edge[0],x], t_data_df.loc[edge[1],x], None]
Yed += [t_data_df.loc[edge[0],y], t_data_df.loc[edge[1],y], None]
Zed += [t_data_df.loc[edge[0],z], t_data_df.loc[edge[1],z], None]
plots.append(go.Scatter3d(x=Xed,
y=Yed,
z=Zed,
mode='lines',
line=go.scatter3d.Line(color='rgb(210,210,210)', width=2),
hoverinfo='none'
))
self.plots.extend(plots)
def plot(self, output_fname, axes_off=False):
"""Plot embedding of patches to html file.
Parameters
----------
output_fname:str
Output html file.
axes_off:bool
Remove axes.
"""
if axes_off:
fig = go.Figure(data=self.plots,layout=go.Layout(scene=dict(xaxis=dict(title='',autorange=True,showgrid=False,zeroline=False,showline=False,ticks='',showticklabels=False),
yaxis=dict(title='',autorange=True,showgrid=False,zeroline=False,showline=False,ticks='',showticklabels=False),
zaxis=dict(title='',autorange=True,showgrid=False,zeroline=False,showline=False,ticks='',showticklabels=False))))
else:
fig = go.Figure(data=self.plots)
py.plot(fig, filename=output_fname, auto_open=False)
def to_pil(arr):
"""Numpy array to pil.
Parameters
----------
arr:array
Numpy array.
Returns
-------
Image
PIL Image.
"""
return Image.fromarray(arr.astype('uint8'), 'RGB')
def blend(arr1, arr2, alpha=0.5):
"""Blend 2 arrays together, mixing with alpha.
Parameters
----------
arr1:array
Image 1.
arr2:array
Image 2.
alpha:float
Higher alpha makes image more like image 1.
Returns
-------
array
Resulting image.
"""
return alpha*arr1 + (1.-alpha)*arr2
def prob2rbg(prob, palette, arr):
"""Convert probability score to rgb image.
Parameters
----------
prob:float
Between 0 and 1 score.
palette:palette
Pallet converts between prob and color.
arr:array
Original array.
Returns
-------
array
New image colored by prediction score.
"""
col = palette(prob)
for i in range(3):
arr[...,i] = int(col[i]*255)
return arr
def seg2rgb(seg, palette, n_segmentation_classes):
"""Color each pixel by segmentation class.
Parameters
----------
seg:array
Segmentation mask.
palette:palette
Color to RGB map.
n_segmentation_classes:int
Total number segmentation classes.
Returns
-------
array
Returned segmentation image.
"""
#print(seg.shape)
#print((seg/n_segmentation_classes))
img=(palette(seg/n_segmentation_classes)[...,:3]*255).astype(int)
#print(img.shape)
return img
def annotation2rgb(i,palette,arr):
"""Go from annotation of patch to color.
Parameters
----------
i:int
Annotation index.
palette:palette
Index to color mapping.
arr:array
Image array.
Returns
-------
array
Resulting image.
"""
col = palette[i]
for i in range(3):
arr[...,i] = int(col[i]*255)
return arr
def plot_image_(image_file, compression_factor=2., test_image_name='test.png'):
"""Plots entire SVS/other image.
Parameters
----------
image_file:str
Image file.
compression_factor:float
Amount to shrink each dimension of image.
test_image_name:str
Output image file.
"""
from pathflowai.utils import svs2dask_array, npy2da
import cv2
if image_file.endswith('.zarr'):
arr=da.from_zarr(image_file)
else:
arr=svs2dask_array(image_file, tile_size=1000, overlap=0, remove_last=True, allow_unknown_chunksizes=False) if (not image_file.endswith('.npy')) else npy2da(image_file)
arr2=to_pil(cv2.resize(arr.compute(), dsize=tuple((np.array(arr.shape[:2])/compression_factor).astype(int).tolist()), interpolation=cv2.INTER_CUBIC))
arr2.save(test_image_name)
# for now binary output
class PredictionPlotter:
"""Plots predictions over entire image.
Parameters
----------
dask_arr_dict:dict
Stores all dask arrays corresponding to all of the images.
patch_info_db:str
Patch level information, eg. prediction.
compression_factor:float
How much to compress image by.
alpha:float
Low value assigns higher weight to prediction over original image.
patch_size:int
Patch size.
no_db:bool
Don't use patch information.
plot_annotation:bool
Plot annotations from patch information.
segmentation:bool
Plot segmentation mask.
n_segmentation_classes:int
Number segmentation classes.
input_dir:str
Input directory.
annotation_col:str
Annotation column to plot.
scaling_factor:float
Multiplies the prediction scores to make them appear darker on the images when predicting.
"""
# some patches have been filtered out, not one to one!!! figure out
def __init__(self, dask_arr_dict, patch_info_db, compression_factor=3, alpha=0.5, patch_size=224, no_db=False, plot_annotation=False, segmentation=False, n_segmentation_classes=4, input_dir='', annotation_col='annotation', scaling_factor=1.):
self.segmentation = segmentation
self.scaling_factor=scaling_factor
self.segmentation_maps = None
self.n_segmentation_classes=float(n_segmentation_classes)
self.pred_palette = sns.cubehelix_palette(start=0,as_cmap=True)
if not no_db:
self.compression_factor=compression_factor
self.alpha = alpha
self.patch_size = patch_size
conn = sqlite3.connect(patch_info_db)
patch_info=pd.read_sql('select * from "{}";'.format(patch_size),con=conn)
conn.close()
self.annotations = {str(a):i for i,a in enumerate(patch_info['annotation'].unique().tolist())}
self.plot_annotation=plot_annotation
self.palette=sns.color_palette(n_colors=len(list(self.annotations.keys())))
#print(self.palette)
if 'y_pred' not in patch_info.columns:
patch_info['y_pred'] = 0.
self.patch_info=patch_info[['ID','x','y','patch_size','annotation',annotation_col]] # y_pred
if 0:
for ID in predictions:
patch_info.loc[patch_info["ID"]==ID,'y_pred'] = predictions[ID]
self.patch_info = self.patch_info[np.isin(self.patch_info['ID'],np.array(list(dask_arr_dict.keys())))]
if self.segmentation:
self.segmentation_maps = {slide:npy2da(join(input_dir,'{}_mask.npy'.format(slide))) for slide in dask_arr_dict.keys()}
#self.patch_info[['x','y','patch_size']]/=self.compression_factor
self.dask_arr_dict = {k:v[...,:3] for k,v in dask_arr_dict.items()}
def add_custom_segmentation(self, basename, npy):
"""Replace segmentation mask with new custom segmentation.
Parameters
----------
basename:str
Patient ID
npy:str
Numpy mask.
"""
self.segmentation_maps[basename] = da.from_array(np.load(npy,mmap_mode='r+'))
def generate_image(self, ID):
"""Generate the image array for the whole slide image with predictions overlaid.
Parameters
----------
ID:str
patient ID.
Returns
-------
array
Resulting overlaid whole slide image.
"""
patch_info = self.patch_info[self.patch_info['ID']==ID]
dask_arr = self.dask_arr_dict[ID]
arr_shape = np.array(dask_arr.shape).astype(float)
#image=da.zeros_like(dask_arr)
arr_shape[:2]/=self.compression_factor
arr_shape=arr_shape.astype(int).tolist()
img = Image.new('RGB',arr_shape[:2],'white')
for i in range(patch_info.shape[0]):
ID,x,y,patch_size,annotation,pred = patch_info.iloc[i].tolist()
#print(x,y,annotation)
x_new,y_new = int(x/self.compression_factor),int(y/self.compression_factor)
image = np.zeros((patch_size,patch_size,3))
if self.segmentation:
image=seg2rgb(self.segmentation_maps[ID][x:x+patch_size,y:y+patch_size].compute(),self.pred_palette, self.n_segmentation_classes)
else:
image=prob2rbg(pred*self.scaling_factor, self.pred_palette, image) if not self.plot_annotation else annotation2rgb(self.annotations[str(pred)],self.palette,image) # annotation
arr=dask_arr[x:x+patch_size,y:y+patch_size].compute()
#print(image.shape)
blended_patch=blend(arr,image, self.alpha).transpose((1,0,2))
blended_patch_pil = to_pil(blended_patch)
patch_size/=self.compression_factor
patch_size=int(patch_size)
blended_patch_pil=blended_patch_pil.resize((patch_size,patch_size))
img.paste(blended_patch_pil, box=(x_new,y_new), mask=None)
return img
def return_patch(self, ID, x, y, patch_size):
"""Return one single patch instead of entire image.
Parameters
----------
ID:str
Patient ID
x:int
X coordinate.
y:int
Y coordinate.
patch_size:int
Patch size.
Returns
-------
array
Image.
"""
img=(self.dask_arr_dict[ID][x:x+patch_size,y:y+patch_size].compute() if not self.segmentation else seg2rgb(self.segmentation_maps[ID][x:x+patch_size,y:y+patch_size].compute(),self.pred_palette, self.n_segmentation_classes))
return to_pil(img)
def output_image(self, img, filename, tif=False):
"""Output calculated image to file.
Parameters
----------
img:array
Image.
filename:str
Output file name.
tif:bool
Store in TIF format?
"""
if tif:
from tifffile import imwrite
imwrite(filename, np.array(img), photometric='rgb')
else:
img.save(filename)
def plot_shap(model, dataset_opts, transform_opts, batch_size, outputfilename, n_outputs=1, method='deep', local_smoothing=0.0, n_samples=20, pred_out=False):
"""Plot shapley attributions overlaid on images for classification tasks.
Parameters
----------
model:nn.Module
Pytorch model.
dataset_opts:dict
Options used to configure dataset
transform_opts:dict
Options used to configure transformers.
batch_size:int
Batch size for training.
outputfilename:str
Output filename.
n_outputs:int
Number of top outputs.
method:str
Gradient or deep explainer.
local_smoothing:float
How much to smooth shapley map.
n_samples:int
Number shapley samples to draw.
pred_out:bool
Label images with binary prediction score?
"""
import torch
from torch.nn import functional as F
import numpy as np
from torch.utils.data import DataLoader
import shap
from pathflowai.datasets import DynamicImageDataset
import matplotlib
from matplotlib import pyplot as plt
from pathflowai.sampler import ImbalancedDatasetSampler
out_transform=dict(sigmoid=F.sigmoid,softmax=F.softmax,none=lambda x: x)
binary_threshold=dataset_opts.pop('binary_threshold')
num_targets=dataset_opts.pop('num_targets')
dataset = DynamicImageDataset(**dataset_opts)
if dataset_opts['classify_annotations']:
binarizer=dataset.binarize_annotations(num_targets=num_targets,binary_threshold=binary_threshold)
num_targets=len(dataset.targets)
dataloader_val = DataLoader(dataset,batch_size=batch_size, num_workers=10, shuffle=True if num_targets>1 else False, sampler=ImbalancedDatasetSampler(dataset) if num_targets==1 else None)
#dataloader_test = DataLoader(dataset,batch_size=batch_size,num_workers=10, shuffle=False)
background,y_background=next(iter(dataloader_val))
if method=='gradient':
background=torch.cat([background,next(iter(dataloader_val))[0]],0)
X_test,y_test=next(iter(dataloader_val))
if torch.cuda.is_available():
background=background.cuda()
X_test=X_test.cuda()
if pred_out!='none':
if torch.cuda.is_available():
model2=model.cuda()
y_test=out_transform[pred_out](model2(X_test)).detach().cpu()
y_test=y_test.numpy()
if method=='deep':
e = shap.DeepExplainer(model, background)
s=e.shap_values(X_test, ranked_outputs=n_outputs)
elif method=='gradient':
e = shap.GradientExplainer(model, background, batch_size=batch_size, local_smoothing=local_smoothing)
s=e.shap_values(X_test, ranked_outputs=n_outputs, nsamples=n_samples)
if y_test.shape[1]>1:
y_test=y_test.argmax(axis=1)
if n_outputs>1:
shap_values, idx = s
else:
shap_values, idx = s, y_test
#print(shap_values) # .detach().cpu()
if num_targets == 1:
shap_numpy = [np.swapaxes(np.swapaxes(shap_values, 1, -1), 1, 2)]
else:
shap_numpy = [np.swapaxes(np.swapaxes(s, 1, -1), 1, 2) for s in shap_values]
#print(shap_numpy.shape)
X_test_numpy=X_test.detach().cpu().numpy()
X_test_numpy=X_test_numpy.transpose((0,2,3,1))
for i in range(X_test_numpy.shape[0]):
X_test_numpy[i,...]*=np.array(transform_opts['std'])
X_test_numpy[i,...]+=np.array(transform_opts['mean'])
X_test_numpy=X_test_numpy.transpose((0,3,1,2))
test_numpy = np.swapaxes(np.swapaxes(X_test_numpy, 1, -1), 1, 2)
if pred_out!='none':
labels=y_test.astype(str)
else:
labels = np.array([[(dataloader_val.dataset.targets[i[j]] if num_targets>1 else str(i)) for j in range(n_outputs)] for i in idx])#[:,np.newaxis] # y_test
if 0 and (len(labels.shape)<2 or labels.shape[1]==1):
labels=labels.flatten()#[:np.newaxis]
#print(labels.shape,shap_numpy.shape[0])
plt.figure()
shap.image_plot(shap_numpy, test_numpy, labels)# if num_targets!=1 else shap_values -test_numpy , labels=dataloader_test.dataset.targets)
plt.savefig(outputfilename, dpi=300)
def plot_umap_images(dask_arr_dict, embeddings_file, ID=None, cval=1., image_res=300., outputfname='output_embedding.png', mpl_scatter=True, remove_background_annotation='', max_background_area=0.01, zoom=0.05, n_neighbors=10, sort_col='', sort_mode='asc'):
"""Make UMAP embedding plot, overlaid with images.
Parameters
----------
dask_arr_dict:dict
Stored dask arrays for each WSI.
embeddings_file:str
Embeddings pickle file stored from running using after trainign the model.
ID:str
Patient ID.
cval:float
Deprecated
image_res:float
Image resolution.
outputfname:str
Output image file.
mpl_scatter:bool
Recommended: Use matplotlib for scatter plot.
remove_background_annotation:str
Remove the background annotations. Enter for annotation to remove.
max_background_area:float
Maximum backgrund area in each tile for inclusion.
zoom:float
How much to zoom in on each patch, less than 1 is zoom out.
n_neighbors:int
Number of neighbors for UMAP embedding.
sort_col:str
Patch info column to sort on.
sort_mode:str
Sort ascending or descending.
Returns
-------
type
Description of returned object.
Inspired by: https://gist.github.com/lukemetz/be6123c7ee3b366e333a
WIP!! Needs testing."""
import torch
import dask
from dask.distributed import Client
from umap import UMAP
from pathflowai.visualize import PlotlyPlot
import pandas as pd, numpy as np
import skimage.io
from skimage.transform import resize
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
sns.set(style='white')
def min_resize(img, size):
"""
Resize an image so that it is size along the minimum spatial dimension.
"""
w, h = map(float, img.shape[:2])
if min([w, h]) != size:
if w <= h:
img = resize(img, (int(round((h/w)*size)), int(size)))
else:
img = resize(img, (int(size), int(round((w/h)*size))))
return img
#dask_arr = dask_arr_dict[ID]
embeddings_dict=torch.load(embeddings_file)
embeddings=embeddings_dict['embeddings']
patch_info=embeddings_dict['patch_info']
if sort_col:
idx=np.argsort(patch_info[sort_col].values)
if sort_mode == 'desc':
idx=idx[::-1]
patch_info = patch_info.iloc[idx]
embeddings=embeddings.iloc[idx]
if ID:
removal_bool=(patch_info['ID']==ID).values
patch_info = patch_info.loc[removal_bool]
embeddings=embeddings.loc[removal_bool]
if remove_background_annotation:
removal_bool=(patch_info[remove_background_annotation]<=(1.-max_background_area)).values
patch_info=patch_info.loc[removal_bool]
embeddings=embeddings.loc[removal_bool]
umap=UMAP(n_components=2,n_neighbors=n_neighbors)
t_data=pd.DataFrame(umap.fit_transform(embeddings.iloc[:,:-1].values),columns=['x','y'],index=embeddings.index)
images=[]
for i in range(patch_info.shape[0]):
ID=patch_info.iloc[i]['ID']
x,y,patch_size=patch_info.iloc[i][['x','y','patch_size']].values.tolist()
arr=dask_arr_dict[ID][x:x+patch_size,y:y+patch_size]#.transpose((2,0,1))
images.append(arr)
c=Client()
images=dask.compute(images)
c.close()
if mpl_scatter:
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
def imscatter(x, y, ax, imageData, zoom):
images = []
for i in range(len(x)):
x0, y0 = x[i], y[i]
img = imageData[i]
#print(img.shape)
image = OffsetImage(img, zoom=zoom)
ab = AnnotationBbox(image, (x0, y0), xycoords='data', frameon=False)
images.append(ax.add_artist(ab))
ax.update_datalim(np.column_stack([x, y]))
ax.autoscale()
fig, ax = plt.subplots()
imscatter(t_data['x'].values, t_data['y'].values, imageData=images[0], ax=ax, zoom=zoom)
sns.despine()
plt.savefig(outputfname,dpi=300)
else:
xx=t_data.iloc[:,0]
yy=t_data.iloc[:,1]
images = [min_resize(image, img_res) for image in images]
max_width = max([image.shape[0] for image in images])
max_height = max([image.shape[1] for image in images])
x_min, x_max = xx.min(), xx.max()
y_min, y_max = yy.min(), yy.max()
# Fix the ratios
sx = (x_max-x_min)
sy = (y_max-y_min)
if sx > sy:
res_x = sx/float(sy)*res
res_y = res
else:
res_x = res
res_y = sy/float(sx)*res
canvas = np.ones((res_x+max_width, res_y+max_height, 3))*cval
x_coords = np.linspace(x_min, x_max, res_x)
y_coords = np.linspace(y_min, y_max, res_y)
for x, y, image in zip(xx, yy, images):
w, h = image.shape[:2]
x_idx = np.argmin((x - x_coords)**2)
y_idx = np.argmin((y - y_coords)**2)
canvas[x_idx:x_idx+w, y_idx:y_idx+h] = image
skimage.io.imsave(outputfname, canvas)
Datasets
Models
