normalize all images in a directory by the percentile of the estimated foreground

normalize-percentile-foreground.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Normalize the intensity of a set of images by 
finding the specified percentile of each image
foreground and voxel-wise dividing the image by that value
Author: Jacob Reinhold
"""

import os
import re
import sys
import warnings
from argparse import ArgumentParser
from pathlib import Path
from typing import List, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchio as tio
from scipy.ndimage.filters import gaussian_filter1d


def plot(
    data: List[np.ndarray],
    out_filename: str,
    title: Optional[str] = None,
    n_bins: int = 200,
    alpha: float = 0.5,
    lw: float = 3.0,
    log: bool = True,
    smooth: bool = True,
    figsize: Tuple[int, int] = (8, 8),
) -> None:
    fig, ax = plt.subplots(figsize=figsize)
    for datum in data:
        hist, bin_edges = np.histogram(datum.flatten(), n_bins)
        bins = np.diff(bin_edges) / 2 + bin_edges[:-1]
        if log:
            # catch divide by zero warnings in call to log
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                hist = np.log10(hist)
                hist[np.isinf(hist)] = 0.0
        if smooth:
            hist = gaussian_filter1d(hist, sigma=1.0)
        ax.plot(bins, hist, alpha=alpha, linewidth=lw)
    ax.set_xlabel("Intensity")
    ax.set_ylabel(r"Log$_{10}$ Count")
    ax.set_ylim((0, None))
    if title is not None:
        ax.set_title(title)
    plt.savefig(out_filename)


def main() -> int:
    parser = ArgumentParser(description="Normalize images for image processing")
    parser.add_argument("top_dir", type=Path)
    parser.add_argument("out_dir", type=Path)
    parser.add_argument("-i", "--include", type=str, default=[".*"], nargs="+")
    parser.add_argument("-e", "--exclude", type=str, default=[None], nargs="+")
    parser.add_argument("-p", "--percentile", type=float, default=50.0)
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("-d", "--dry-run", action="store_true")
    parser.add_argument("--plot", action="store_true")
    args = parser.parse_args()
    include_progs = [re.compile(inc) for inc in args.include]
    exclude_progs = [re.compile(exc) for exc in args.exclude if exc is not None]
    percentiles = []
    foregrounds_before = []
    foregrounds_after = []
    try:
        for root, dirs, files in os.walk(args.top_dir):
            for _fn in files:
                if all(inc.search(_fn) is None for inc in include_progs) or any(
                    exc.search(_fn) is not None for exc in exclude_progs
                ):
                    if args.verbose:
                        print(f"Skipping: {_fn}")
                    continue
                fn = Path(root) / _fn
                if args.verbose:
                    print(f"Normalizing: {str(fn)}")
                out_fn = str(fn).replace(str(fn.parents[len(fn.parents) - 2]) + "/", "")
                out_path = args.out_dir / out_fn
                image = tio.ScalarImage(fn)
                if not args.dry_run:
                    img = image.numpy()
                    mask = img > img.mean()
                    foreground = img[mask]
                    if args.plot:
                        foregrounds_before.append(foreground)
                    percentile = np.percentile(foreground, args.percentile)
                    percentiles.append(percentile)
                else:
                    percentile = np.random.randn()
                    percentiles.append(percentile)
                if args.verbose:
                    print(f"Saving normalized: {str(out_path)}; p={percentile:0.3e}")
                if not args.dry_run:
                    out_path.parent.mkdir(parents=True, exist_ok=True)
                    normalized = img / percentile
                    if args.plot:
                        foregrounds_after.append(normalized[mask])
                    image.set_data(torch.from_numpy(normalized))
                    image.save(out_path)
    except KeyboardInterrupt:
        ...
    percentile = np.median(percentiles)
    std = np.std(percentiles)
    if args.verbose:
        print(f"Median: {percentile}; Std: {std}")
    if args.plot and not args.dry_run:
        plot(foregrounds_before, "before.png", title="Before normalization")
        plot(foregrounds_after, "after.png", title="After normalization")
    return 0


if __name__ == "__main__":
    sys.exit(main())

If you use this script in an academic paper, please cite the paper:

    @inproceedings{reinhold2019evaluating,
      title={Evaluating the impact of intensity normalization on {MR} image synthesis},
      author={Reinhold, Jacob C and Dewey, Blake E and Carass, Aaron and Prince, Jerry L},
      booktitle={Medical Imaging 2019: Image Processing},
      volume={10949},
      pages={109493H},
      year={2019},
      organization={International Society for Optics and Photonics}}