Notebook to predict the segmentation masks on new data using pretrained or customized models.
#@markdown Please run this cell to get started.
%load_ext autoreload
%autoreload 2
try:
from google.colab import files, drive
except ImportError:
pass
try:
import deepflash2
except ImportError:
!pip install -q deepflash2==0.0.14
import zipfile
import shutil
import imageio
from fastai.vision.all import *
from deepflash2.all import *
from scipy.stats import entropy
Required data structure: One folder for images (different from training image folder!)
Examplary structure:
[folder] images_new
- [file] 0001.tif
- [file] 0002.tif
- The folder in your drive must contain all images in one folder.
- See here how to organize your files in Google Drive.
- See this stackoverflow post for browsing files with the file browser
try:
drive.mount('/content/drive')
path = "/content/drive/My Drive/data" #@param {type:"string"}
path = Path(path)
print('Path contains the following files and folders: \n', L(os.listdir(path)))
#@markdown Follow the instructions and press Enter after copying and pasting the key.
except:
print("Warning: Connecting to Google Drive only works on Google Colab.")
pass
- The zip file must contain all images in one folder
- See here how to zip files on Windows or Mac.
path = Path('data')
try:
u_dict = files.upload()
for key in u_dict.keys():
unzip(path, key)
print('Path contains the following files and folders: \n', L(os.listdir(path)))
except:
print("Warning: File upload only works on Google Colab.")
pass
If you're working on your local machine or server, provide a path to the correct folder.
path = "data" #@param {type:"string"}
path = Path(path)
print('Path contains the following files and folders: \n', L(os.listdir(path)))
If you don't have any data available yet, try our sample data [Description]
path = Path('sample_data_cFOS')
url = "https://github.com/matjesg/deepflash2/releases/download/model_library/wue1_cFOS_small.zip"
urllib.request.urlretrieve(url, 'sample_data_cFOS.zip')
unzip(path, 'sample_data_cFOS.zip')
_ = shutil.move(path/'images', path/'images_new')
image_folder = "images_new" #@param {type:"string"}
f_names = get_image_files(path/image_folder)
print(f'Found {len(f_names)} images in "{path}".')
#@markdown Number of classes: e.g., 2 for binary segmentation (foreground and background class)
n_classes = 2 #@param {type:"integer"}
ds = TileDataset(f_names, n_classes=n_classes)
#@markdown Use the slider to control the number of displayed images
first_n = 3 #@param {type:"slider", min:1, max:100, step:1}
ds.show_data(max_n=first_n, figsize=(5,5), overlay=False)
- Select the same model architecture as used during Model Training
model_arch = 'unet_deepflash2' #@param ["unet_deepflash2", "unet_falk2019", "unet_ronnberger2015"]
n_channels = ds.get_data(max_n=1)[0].shape[-1]
model = torch.hub.load('matjesg/deepflash2', model_arch, pretrained=False, n_classes=ds.c, in_channels=n_channels, force_reload=True)
model_list = L()
- Works with Google Drive connection or local installation.
- If you haven't already done it, train your own models in the Model Training Notebook
#@markdown Run to select models from `models_folder`.
#@markdown Models should be saved in the 'models' folder of your provided path.
models_folder = "models" #@param {type:"string"}
model_list = get_files(path/models_folder, extensions='.pth')
print('Found models', model_list)
try:
u_dict = files.upload()
model_list += [Path(u) for u in u_dict]
except:
print("Warning: File upload only works on Google Colab.")
pass
print('Found models', model_list)
- Out-of-the-box prediction without prior validation can lead to invalid results.
- See Model Libray for more information.
pretrained_weights = "wue_cFOS" #@param ["wue_cFOS", "wue_Parv", "wue_GFAP", "wue_GFP", "wue_OPN3"]
model = torch.hub.load('matjesg/deepflash2', model_arch, pretrained=True, dataset=pretrained_weights, n_classes=ds.c, in_channels=n_channels)
model_list = L(Path(f'{torch.hub.get_dir()}/checkpoints/{pretrained_weights}.pth'))
res, res_mc = {}, {}
for m in progress_bar(model_list):
print(f'Model {m.stem}')
dls = DataLoaders.from_dsets(ds, batch_size=4 ,shuffle=False, drop_last=False)
state_dict = torch.load(m)
model.load_state_dict(state_dict, strict=False)
if torch.cuda.is_available(): dls.cuda(), model.cuda()
learn = Learner(dls, model, loss_func=0)#.to_fp16()
print(f'Predicting segmentation masks')
smxs, segs, _ = learn.predict_tiles(dl=dls.train)
print(f'Predicting uncertainty maps')
smxs_mc, segs_mc, std = learn.predict_tiles(dl=dls.train, mc_dropout=True, n_times=10)
#TODO Save results not using RAM
for i, file in enumerate(f_names):
res[(m.stem, file)] = smxs[i], segs[i]
res_mc[(m.stem, file)] = smxs_mc[i], segs_mc[i], std[i]
In this section you can ensemble (combine) the single model outputs.
pred_dir = 'preds' #@param {type:"string"}
pred_path = path/pred_dir/'ensemble'
pred_path.mkdir(parents=True, exist_ok=True)
uncertainty_dir = 'uncertainties' #@param {type:"string"}
uncertainty_path = path/uncertainty_dir/'ensemble'
uncertainty_path.mkdir(parents=True, exist_ok=True)
result_path = path/'results'
result_path.mkdir(exist_ok=True)
#@markdown Define `filetype` to save the predictions and uncertainties. All common filetypes are supported.
filetype = 'png' #@param {type:"string"}
res_list = []
for file in f_names:
img = ds.get_data(file)[0]
pred = ensemble_results(res, file)
pred_std = ensemble_results(res_mc, file, std=True)
df_tmp = pd.Series({'file' : file.name, 'entropy': entropy(pred_std, axis=None)})
plot_results(img, pred, pred_std, df=df_tmp)
res_list.append(df_tmp)
imageio.imsave(pred_path/f'{file.name}_pred.{filetype}', pred.astype(np.uint8) if np.max(pred)>1 else pred.astype(np.uint8)*255)
imageio.imsave(uncertainty_path/f'{file.name}_uncertainty.{filetype}', pred_std.astype(np.uint8)*255)
df_res = pd.DataFrame(res_list)
df_res.to_csv(result_path/'ensemble_results.csv', index=False)
res_list = []
for file in f_names:
print(f'###### File {file.name} ######')
img = ds.get_data(file)[0]
for model_path in model_list:
model_name = model_path.stem
val_files = [f for mod , f in res.keys() if mod == model_name]
print(f'{model_name}')
pred_path = path/pred_dir/model_name
pred_path.mkdir(parents=True, exist_ok=True)
uncertainty_path = path/uncertainty_dir/model_name
uncertainty_path.mkdir(parents=True, exist_ok=True)
pred = res[(model_name,file)][1]
pred_std = res_mc[(model_name,file)][2][...,0]
df_tmp = pd.Series({'file' : file.name,
'model' : model_name,
'entropy': entropy(pred_std, axis=None)})
plot_results(img, pred, pred_std, df=df_tmp)
res_list.append(df_tmp)
imageio.imsave(pred_path/f'{file.stem}_pred.{filetype}', pred.astype(np.uint8) if np.max(pred)>1 else pred.astype(np.uint8)*255)
imageio.imsave(uncertainty_path/f'{file.stem}_uncertainty.{filetype}', pred_std.astype(np.uint8)*255)
pd.DataFrame(res_list).to_csv(result_path/'model_results.csv', index=False)
To download validation predictions and uncertainties, you first need to execute Section Validate models and ensembles.
Note: If you're connected to Google Drive, the models are automatically saved to your drive.
out_name = 'predictions'
shutil.make_archive(path/out_name, 'zip', path/pred_dir)
try:
files.download(path/f'{out_name}.zip')
except:
print("Warning: File download only works on Google Colab.")
pass
out_name = 'uncertainties'
shutil.make_archive(path/out_name, 'zip', path/uncertainty_dir)
try:
files.download(path/f'{out_name}.zip')
except:
print("Warning: File download only works on Google Colab.")
pass
try:
files.download(result_path/f'ensemble_results.csv')
except:
print("Warning: File download only works on Google Colab.")
pass
try:
files.download(result_path/'model_results.csv')
except:
print("Warning: File download only works on Google Colab.")
pass