Quickly Erode Masks from Warp

warp_mask_eroder.py

# Author: Benjamin Barad
# A quick script to erode the masks generated by warp to more tightly surround frost and beads
# Warp's neural net masking can be used to quickly remove gold beads and hunks of frost from tomograms. 
# However, it tends to overmask and generate a nice big blob around each bead. This is desirable for
# limiting particle picking, but not for tomogram generation. This script erodes the mask to more tightly
# surround the frost and beads.
# Usage: 
# 1. Run warp and generate masks for your tomograms with boxnet3mask
# 2. Close warp and rename the warp "mask" folder to "mask_raw"
# 3. Run this script as follows:
# python warp_mask_eroder.py --input_folder /WAPFOLDER/mask_raw --output_folder /WARPFOLDER/mask --erode_iterations 15
# 4. Reopen warp and confirm that the masks are now more tightly surrounding the beads and frost
# 5. If necessary, repeat with a different erode_iterations parameter to tweak to a tight but not too-tight fit.
# Requirements:
# 1. skimage
# 2. scipy
# 3. tqdm
# 4. click
# 5. numpy
# 6. imagecodecs


import skimage.io as io
import numpy as np
import scipy as sp
import tqdm
import click
import glob
import os

def warp_mask_eroder(input_folder, output_folder, erode_iterations, exclude_list=[], mask_struct = None):
    """This function erodes the mask of each image in the input folder and saves the result in the output folder.
    Parameters
    ----------
    input_folder : str
        Path to the folder containing the images to be processed.
    output_folder : str
        Path to the folder where the processed images will be saved.
    erode_iterations : int
        Number of iterations for the erosion.
    exclude_list : list
        List of images to be excluded from the processing.
    mask_struct : array
        Structuring element for the erosion. If None, a 3x3 square is used."""
    if not os.path.exists(input_folder):
        raise ValueError('Input folder does not exist.') 
    if mask_struct is None:
        mask_struct = np.ones((3,3))
    if not os.path.exists(output_folder):
        print('Output folder does not exist. Creating folder...')
        os.makedirs(output_folder)

    for file in tqdm.tqdm(glob.glob(input_folder + '/*.tif')):
        if file not in exclude_list:
            mask = io.imread(file)
            mask = sp.ndimage.binary_erosion(mask, mask_struct, iterations=erode_iterations, border_value=1)
            io.imsave(output_folder + '/' + file.split('/')[-1], mask.astype(np.float32))

@click.command()
@click.option('--input_folder', help='Path to the folder containing the images to be processed.', required=True)
@click.option('--output_folder', help='Path to the folder where the processed images will be saved.', required=True)
@click.option('--erode_iterations', help='Number of iterations for the erosion.', default=1)
@click.option('--exclude_list', help='List of images to be excluded from the processing.', default=[])
@click.option('--mask_struct', help='Structuring element for the erosion. Options are cross and square.', default=None)
def main(input_folder, output_folder, erode_iterations, exclude_list, mask_struct):
    if mask_struct == 'cross':
        mask_struct = np.array([[0,1,0],[1,1,1],[0,1,0]])
    elif mask_struct == 'square':
        mask_struct = np.ones((3,3))
    else:
        mask_struct = None
    warp_mask_eroder(input_folder, output_folder, erode_iterations, exclude_list, mask_struct)

if __name__ == '__main__':
    main()