ct-segmentation-toolkit 0.2.0__py3-none-any.whl

ct_seg/denoise/train.py

@@ -0,0 +1,385 @@
+"""
+main.py — Training entry point for 2.5D Noise2Inverse (N2I).
+
+Trains a self-supervised U-Net that denoises computed-tomography (CT)
+reconstructions without any ground-truth (clean) images, following the
+Noise2Inverse framework. Two statistically independent sub-reconstructions
+(e.g. reconstructions built from even- vs. odd-numbered projections) supervise
+each other: the network learns to predict one split from the other. Because the
+noise is independent between the splits while the underlying object is shared,
+the optimum of this objective is the denoised signal. The "2.5D" variant feeds
+five adjacent axial slices as input channels so the model can use through-plane
+context while still predicting a single 2D slice.
+
+Inputs  : a YAML config pointing at the reconstruction directory (see
+          baseline_config.yaml) plus the even/odd sub-reconstruction TIFF stacks.
+Outputs : a TrainOutput/ directory (created next to the reconstructions) holding
+          the best checkpoints, a training log, and periodic preview PNGs.
+
+Training procedure:
+    * Warm up using a pixel-wise L1 loss only.
+    * After the LCL warm-up, add a Laplacian Contrast Loss term (loss.py) that
+      steers contrast toward edges to keep them sharp.
+    * Each epoch, keep three "best" checkpoints: lowest validation L1 loss,
+      lowest LCL loss, and highest Laplacian edge score (eval.py).
+
+Usage:
+    python main.py -gpus=0 -config=/path/to/config.yaml
+    (normally launched via train.sh)
+
+Author:  Cameron Renteria <crentb23@gmail.com>
+License: Apache-2.0 (see LICENSE)
+"""
+
+import argparse
+import logging
+import os
+import shutil
+import sys
+import time
+from copy import deepcopy
+
+import numpy as np
+import torch
+import yaml
+from matplotlib import pyplot as plt
+from torch.utils.data import DataLoader
+
+from ct_seg.denoise.data import TomoDatasetTrain
+from ct_seg.denoise.eval import laplacian_score_batch
+from ct_seg.denoise.loss import LCL
+from ct_seg.denoise.model import unet_ns_gn
+from ct_seg.denoise.utils import save2img
+
+
+def count_parameters(model):
+    """Return the number of trainable (gradient-requiring) parameters in `model`."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def main(args):
+
+    # read the YAML file
+    with open(args.config, "r") as file:
+        params = yaml.safe_load(file)
+
+    START_TIME = time.time()
+
+    # create directory containing training results in the directory of the reconstructions
+    path_to_reconstructions = params["dataset"]["directory_to_reconstructions"]
+    odir = path_to_reconstructions + "/" + "TrainOutput"
+    if os.path.isdir(odir):
+        shutil.rmtree(odir)
+    os.mkdir(odir)
+    os.mkdir(f"{odir}/results")
+
+    # create output log
+    logging.basicConfig(filename=f"{odir}/Noise2Inverse.log", level=logging.DEBUG)
+    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+
+    # setup device
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if dev.type == "cuda":
+        torch.cuda.set_device(0)
+    logging.info(f"Using device: {dev}")
+
+    logging.info("\nLoading data into CPU memory, it will take a while ... ...")
+    ds_train = TomoDatasetTrain(params=params, config_file=args.config)
+    dl_train = DataLoader(
+        dataset=ds_train,
+        batch_size=params["train"]["mbsz"],
+        shuffle=True,
+        num_workers=4,
+        drop_last=False,
+        prefetch_factor=2,
+        pin_memory=True,
+    )
+
+    logging.info(
+        f"\nLoaded %d samples, {ds_train.samples}, into CPU memory for training." % (len(ds_train),)
+    )
+
+    # initialized model from scratch
+    model = unet_ns_gn(ich=5, start_filter_size=16, channels_per_group=8).to(dev)
+    optimizer = torch.optim.Adam(model.parameters(), lr=params["train"]["lr"])
+    logging.info("\nInitializing model from scratch\n")
+
+    print(f"Number of model parameters: {count_parameters(model):,}")
+    model_updates = 0
+
+    # Loss functions and warm-up schedule:
+    #   criterion     : pixel-wise L1 loss (used for the entire run)
+    #   criterion_lcl : Laplacian Contrast Loss, the edge-sharpening term (loss.py)
+    #   beta          : weight applied to the LCL term once it is switched on
+    #   lcl_warmup    : number of model updates of pure-L1 warm-up before LCL kicks in
+    criterion = torch.nn.L1Loss()
+    criterion_lcl = LCL()
+    beta = 0.01
+    lcl_warmup = 1500
+    warmup = params["train"]["warmup"]
+    continue_warmup = True
+
+    train_loss, val_loss = [], []
+    edge_values = []
+    train_lcl_loss, val_lcl_loss = [], []
+
+    best_val_loss, best_edge, best_lcl_loss = np.inf, 0, np.inf
+    best_val_epoch, best_edge_epoch, best_lcl_epoch = 0, 0, 0
+
+    # start training
+    for epoch in range(1, params["train"]["maxep"] + 1):
+
+        step_losses, step_val_losses, step_lcl_loss, step_lcl_val_loss, step_edge_values = (
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+
+        tick_ep = time.time()
+
+        model.train()
+        # training loop
+        for X_mb, Y_mb in dl_train:
+            optimizer.zero_grad()
+
+            X_mb_dev = X_mb.to(dev)
+            Y_mb_dev = Y_mb.to(dev)
+
+            # Symmetric Noise2Inverse update. X_mb and Y_mb are the two independent
+            # sub-reconstruction splits. Each step trains the network in BOTH
+            # directions -- predict the center slice (channel index 2) of one split
+            # from the other -- so neither split is privileged as 'input' vs 'target'.
+            # While model_updates <= lcl_warmup only the L1 term is optimized;
+            # afterwards the edge-aware LCL term (loss.py) is added in.
+            if model_updates <= lcl_warmup:
+                optimizer.zero_grad()
+
+                # Process first view
+                pred_view1 = model(X_mb_dev)
+
+                loss_view1 = criterion(pred_view1.squeeze(dim=1), Y_mb_dev[:, 2])
+
+                loss_view1.backward()
+                optimizer.step()
+
+                optimizer.zero_grad()
+
+                # Process second view
+                pred_view2 = model(Y_mb_dev)
+                loss_view2 = criterion(pred_view2.squeeze(dim=1), X_mb_dev[:, 2])
+
+                loss_view2.backward()
+                optimizer.step()
+
+                loss_lcl1 = torch.tensor(0.0)
+                loss_lcl2 = torch.tensor(0.0)
+
+            else:
+                optimizer.zero_grad()
+
+                # Process first view
+                pred_view1 = model(X_mb_dev)
+
+                loss_view1 = criterion(pred_view1.squeeze(dim=1), Y_mb_dev[:, 2])
+                loss_lcl1 = criterion_lcl(pred_view1) * beta
+                loss_total1 = loss_view1 + loss_lcl1
+
+                loss_total1.backward()
+                optimizer.step()
+
+                optimizer.zero_grad()
+
+                # Process second view
+                pred_view2 = model(Y_mb_dev)
+                loss_view2 = criterion(pred_view2.squeeze(dim=1), X_mb_dev[:, 2])
+                loss_lcl2 = criterion_lcl(pred_view2) * beta
+                loss_total2 = loss_view2 + loss_lcl2
+
+                loss_total2.backward()
+                optimizer.step()
+
+            loss = loss_view1 + loss_view2
+            step_losses.append(loss.detach().cpu().numpy())
+            loss_lcl = loss_lcl1 + loss_lcl2
+            step_lcl_loss.append(loss_lcl.detach().cpu().numpy())
+            model_updates += 1
+
+        model.eval()
+        with torch.no_grad():
+            # validation loop
+            for X_mb, Y_mb in dl_train:
+                X_mb_dev = X_mb.to(dev)
+                Y_mb_dev = Y_mb.to(dev)
+
+                pred_view1 = model(X_mb_dev)
+                loss_view1 = criterion(pred_view1.squeeze(dim=1), Y_mb_dev[:, 2])
+                loss_lcl1 = criterion_lcl(pred_view1) * beta
+
+                pred_view2 = model(Y_mb_dev)
+                loss_view2 = criterion(pred_view2.squeeze(dim=1), X_mb_dev[:, 2])
+                loss_lcl2 = criterion_lcl(pred_view2) * beta
+
+                loss = loss_view1 + loss_view2
+
+                lap_score = laplacian_score_batch(pred_view1.cpu()) + laplacian_score_batch(
+                    pred_view2.cpu()
+                )
+                step_edge_values.append(lap_score)
+
+                step_val_losses.append(loss.detach().cpu().numpy())
+                step_lcl_val_loss.append(loss_lcl.cpu().numpy())
+
+        ep_time = time.time() - tick_ep
+        logging.info(f"\nEpoch {epoch}")
+        iter_prints = f"[Train] L1 loss:    {np.mean(step_losses):.6f}, {step_losses[0]:.6f} => {step_losses[-1]:.6f}, rate: {ep_time:.2f}s/ep"
+        logging.info(iter_prints)
+        iter_prints = f"[Train] LCL loss:   {np.mean(step_lcl_loss):.6f}, {step_lcl_loss[0]:.6f} => {step_lcl_loss[-1]:.6f}, rate: {ep_time:.2f}s/ep"
+        logging.info(iter_prints)
+        iter_prints = f"[Val]   L1 loss:    {np.mean(step_val_losses):.6f}, {step_val_losses[0]:.6f} => {step_val_losses[-1]:.6f}, rate: {ep_time:.2f}s/ep"
+        logging.info(iter_prints)
+        iter_prints = f"[Val]   LCL loss:   {np.mean(step_lcl_val_loss):.6f}, {step_lcl_val_loss[0]:.6f} => {step_lcl_val_loss[-1]:.6f}, rate: {ep_time:.2f}s/ep"
+        logging.info(iter_prints)
+        iter_prints = f"[Val] EDGE Value:   {np.mean(step_edge_values):.4f}, {step_edge_values[0]:.4f} => {step_edge_values[-1]:.4f}, rate: {ep_time:.2f}s/ep"
+        logging.info(iter_prints)
+
+        train_loss.append(np.mean(step_losses))
+        val_loss.append(np.mean(step_val_losses))
+        train_lcl_loss.append(np.mean(step_lcl_loss))
+        val_lcl_loss.append(np.mean(step_lcl_val_loss))
+        edge_values.append(np.mean(step_edge_values))
+
+        # Save the best model with the lowest lcl loss
+        if np.mean(step_lcl_val_loss) < best_lcl_loss:
+            best_lcl_loss = np.mean(step_lcl_val_loss)
+            best_lcl_epoch = epoch
+            mdl_fname = f"{odir}/best_lcl_model.pth"
+            torch.save(
+                {
+                    "model_state_dict": deepcopy(model.state_dict()),
+                    "optimizer_state_dict": deepcopy(optimizer.state_dict()),
+                },
+                mdl_fname,
+            )
+
+        # Save the best model with the lowest val loss
+        if np.mean(step_val_losses) < best_val_loss:
+            best_val_loss = np.mean(step_val_losses)
+            best_val_epoch = epoch
+            mdl_fname = f"{odir}/best_val_model.pth"
+            torch.save(
+                {
+                    "model_state_dict": deepcopy(model.state_dict()),
+                    "optimizer_state_dict": deepcopy(optimizer.state_dict()),
+                },
+                mdl_fname,
+            )
+
+        # Save the best model with the highest edge value
+        if np.mean(step_edge_values) > best_edge:
+            best_edge = np.mean(step_edge_values)
+            best_edge_epoch = epoch
+            mdl_fname = f"{odir}/best_edge_model.pth"
+            torch.save(
+                {
+                    "model_state_dict": deepcopy(model.state_dict()),
+                    "optimizer_state_dict": deepcopy(optimizer.state_dict()),
+                },
+                mdl_fname,
+            )
+
+        # End of warm-up: once BOTH warm-up thresholds (warmup and lcl_warmup) are
+        # passed, reset the best-metric trackers so that checkpoints selected during
+        # the pre-LCL phase do not dominate the rest of training.
+        if model_updates > warmup and model_updates > lcl_warmup and continue_warmup:
+
+            best_edge, best_lcl_loss = 0, np.inf
+            best_edge_epoch, best_lcl_epoch = 0, 0
+
+            continue_warmup = False
+
+        CRNT_TIME = time.time()
+        logging.info(f"[Info]  Training Time: {CRNT_TIME-START_TIME:.2f} seconds")
+
+        # option to view the denoising process during training
+        if epoch % 5 == 0:
+            ridx = np.random.randint(pred_view1.shape[0])
+
+            save2img(
+                pred_view1[ridx, -1].detach().cpu().numpy(),
+                "%s/results/_%d_pred_view1.png" % (odir, epoch),
+            )
+            save2img(
+                pred_view2[ridx, -1].detach().cpu().numpy(),
+                "%s/results/_%d_pred_view2.png" % (odir, epoch),
+            )
+            save2img(
+                X_mb_dev[ridx, 2].detach().cpu().numpy(), "%s/results/_%d_gt.png" % (odir, epoch)
+            )
+            save2img(
+                Y_mb_dev[ridx, 2].detach().cpu().numpy(), "%s/results/_%d_noise.png" % (odir, epoch)
+            )
+
+        # Keep track of when/where the best model is
+        logging.info(f"Lowest model validation loss {best_val_loss:.6f} at epoch {best_val_epoch}")
+        logging.info(f"Lowest model LCL loss {best_lcl_loss:.6f} at epoch {best_lcl_epoch}")
+        logging.info(f"Highest model EDGE score {best_edge:.6f} at epoch {best_edge_epoch}")
+        logging.info(f"Number of model updates: {model_updates:,}")
+        logging.info(f"Is model warming up?: {continue_warmup}")
+
+        # View the training/validation loss during training
+        if epoch % 5 == 0:
+            plt.figure(figsize=(12, 8))
+            plt.title("Training Progress")
+            plt.plot(train_loss[:], label="Training Loss")
+            plt.plot(val_loss[:], label="Validation Loss")
+            plt.xlabel("Epoch")
+            plt.ylabel("Loss")
+            plt.legend()
+            plt.savefig(f"{odir}/results/__model_training.png")
+            plt.close()
+
+            plt.figure(figsize=(12, 8))
+            plt.title("Training Progress")
+            plt.plot(train_lcl_loss[25:], label="Training Loss")
+            plt.plot(val_lcl_loss[warmup:], label="Validation Loss")
+            plt.xlabel("Epoch")
+            plt.ylabel("Loss")
+            plt.legend()
+            plt.savefig(f"{odir}/results/__model_lcl_training.png")
+            plt.close()
+
+            plt.figure(figsize=(12, 8))
+            plt.title("Training Progress")
+            plt.plot(edge_values[25:])
+            plt.xlabel("Epoch")
+            plt.ylabel("EDGE Gradient")
+            plt.savefig(f"{odir}/results/__edge_training.png")
+            plt.close()
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Noise2Inverse with 2.5D")
+    parser.add_argument("-gpus", type=str, default="0", help="list of visiable GPUs")
+    parser.add_argument(
+        "-verbose", type=int, default=1, help="1:print to terminal; 0: redirect to file"
+    )
+    parser.add_argument("-config", type=str, required=True, help="path to config yaml file")
+
+    args, unparsed = parser.parse_known_args()
+
+    if len(unparsed) > 0:
+        print("Unrecognized argument(s): \n%s \nProgram exiting ... ..." % "\n".join(unparsed))
+        exit(0)
+
+    if len(args.gpus) > 0:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+
+    logical_cpus = os.cpu_count()
+    os.environ["OMP_NUM_THREADS"] = str(logical_cpus)
+
+    logging.getLogger("matplotlib.font_manager").disabled = True
+
+    main(args)

ct_seg/denoise/utils.py

@@ -0,0 +1,68 @@
+"""
+utils.py — Small image-saving and command-line helper utilities.
+
+Helpers shared across the project: write an array to a TIFF or 8-bit PNG preview
+(`save2img`), save an RGB preview (`save2img_rgb`), rescale an array to uint8
+(`scale2uint8`), and parse boolean command-line strings (`str2bool`).
+
+Author:  Cameron Renteria <crentb23@gmail.com>
+License: Apache-2.0 (see LICENSE)
+"""
+
+import argparse
+
+import skimage.io
+from matplotlib import pyplot as plt
+
+
+def save2img_rgb(img_data, img_fn):
+    """Save an array as an RGB PNG preview, sized to the data on a black background."""
+    plt.figure(figsize=(img_data.shape[1] / 10.0, img_data.shape[0] / 10.0))
+    plt.axes([0, 0, 1, 1])
+    plt.imshow(
+        img_data,
+    )
+    plt.axis("off")
+    plt.savefig(img_fn, facecolor="black", edgecolor="black", dpi=10)
+    plt.close()
+
+
+def save2img(d_img, fn):
+    """Save an array to disk.
+
+    If `fn` ends in 'tiff' the raw values are written unchanged; otherwise the data
+    is min-max scaled to 0-255 and saved as an 8-bit image (e.g. a PNG preview).
+    """
+    if fn[-4:] == "tiff":
+        img_norm = d_img.copy()
+    else:
+        _min, _max = d_img.min(), d_img.max()
+        if _max == _min:
+            img_norm = d_img - _max
+        else:
+            img_norm = (d_img - _min) * 255.0 / (_max - _min)
+        img_norm = img_norm.astype("uint8")
+    skimage.io.imsave(fn, img_norm, check_contrast=False)
+
+
+def scale2uint8(_img):
+    """Min-max scale an array to the 0-255 uint8 range (a constant image maps to 0)."""
+    _min, _max = _img.min(), _img.max()
+    if _max == _min:
+        _img_s = _img - _max
+    else:
+        _img_s = (_img - _min) * 255.0 / (_max - _min)
+    _img_s = _img_s.astype("uint8")
+    return _img_s
+
+
+def str2bool(v):
+    """Parse a truthy/falsy command-line string into a bool (for use with argparse)."""
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")

ct_seg/labeling.py

@@ -0,0 +1,145 @@
+"""
+Napari-based Labeling Tool for CT Segmentation
+===============================================
+Opens a tiff stack in Napari with a labels layer for manual annotation.
+Supports loading existing masks, painting labels, and saving.
+
+Usage:
+    python label_tool.py --images /path/to/tiffs
+    python label_tool.py --images /path/to/tiffs --masks /path/to/existing_masks
+    python label_tool.py --images /path/to/tiffs --num_classes 4 --slice_range 500 600
+
+Controls in Napari:
+    - Select the "Labels" layer in the layer list
+    - Use the paint brush (press B) to draw labels
+    - Use the eraser (press E) to erase
+    - Number keys (1, 2, 3, ...) select the label class
+    - Scroll through slices with the slider at the bottom
+    - Press Ctrl+S or use File > Save to save labels
+
+When you close Napari, masks are automatically saved.
+
+Requirements:
+    pip install napari[all]
+"""
+
+import argparse
+import re
+import sys
+from pathlib import Path
+
+import numpy as np
+import tifffile
+
+
+def natural_sorted(paths):
+    def key(x):
+        return [int(c) if c.isdigit() else c for c in re.split(r"([0-9]+)", str(x))]
+
+    return sorted(paths, key=key)
+
+
+def load_stack(dir_path, start=None, end=None):
+    dir_path = Path(dir_path)
+    paths = natural_sorted(list(dir_path.glob("*.tif*")))
+    if start is not None and end is not None:
+        paths = paths[start:end]
+    if not paths:
+        raise FileNotFoundError(f"No tiff files in {dir_path}")
+    from tqdm import tqdm
+
+    imgs = np.stack([tifffile.imread(str(p)) for p in tqdm(paths, desc="Loading")])
+    return imgs, paths
+
+
+def save_masks(masks, out_dir, offset=0):
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    from tqdm import tqdm
+
+    for i in tqdm(range(masks.shape[0]), desc="Saving masks"):
+        tifffile.imwrite(str(out_dir / f"{i + offset:05d}.tiff"), masks[i].astype(np.uint8))
+    print(f"Saved {masks.shape[0]} masks to {out_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Napari Labeling Tool for CT Segmentation")
+    parser.add_argument("--images", required=True, help="Directory of image tiffs")
+    parser.add_argument("--masks", default=None, help="Directory of existing mask tiffs (optional)")
+    parser.add_argument(
+        "--output", default=None, help="Output directory for masks (default: images_masks)"
+    )
+    parser.add_argument("--num_classes", type=int, default=4, help="Number of classes (default: 4)")
+    parser.add_argument("--slice_range", nargs=2, type=int, default=None, metavar=("START", "END"))
+    args = parser.parse_args()
+
+    try:
+        import napari
+    except ImportError:
+        print("Napari is required for the labeling tool.")
+        print("Install with: pip install 'napari[all]'")
+        sys.exit(1)
+
+    # Setup output
+    if args.output is None:
+        args.output = str(Path(args.images).parent / (Path(args.images).name + "_masks"))
+
+    # Load images
+    start = args.slice_range[0] if args.slice_range else None
+    end = args.slice_range[1] if args.slice_range else None
+    images, paths = load_stack(args.images, start=start, end=end)
+    print(f"Loaded {images.shape[0]} images of size {images.shape[1]}x{images.shape[2]}")
+
+    # Load or create masks
+    if args.masks:
+        masks, _ = load_stack(args.masks, start=start, end=end)
+        print(f"Loaded existing masks with classes: {np.unique(masks)}")
+    else:
+        masks = np.zeros(images.shape, dtype=np.uint8)
+        print("Created empty mask volume")
+
+    # Normalize images for display
+    p2, p98 = np.percentile(images, (2, 98))
+    display_imgs = np.clip(images, p2, p98)
+
+    # Color map for labels
+    label_colors = {
+        0: [0, 0, 0, 0],  # transparent background
+        1: [1, 0.2, 0.2, 0.6],  # red
+        2: [0.2, 1, 0.2, 0.6],  # green
+        3: [0.2, 0.4, 1, 0.6],  # blue
+        4: [1, 1, 0.2, 0.6],  # yellow
+        5: [1, 0.2, 1, 0.6],  # magenta
+        6: [0.2, 1, 1, 0.6],  # cyan
+    }
+
+    # Launch Napari
+    viewer = napari.Viewer(title="CT Segmentation Labeling Tool")
+    viewer.add_image(display_imgs, name="CT Image", colormap="gray")
+    labels_layer = viewer.add_labels(masks, name="Labels", color=label_colors)
+
+    # Set brush size
+    labels_layer.brush_size = 10
+
+    print(f"\n{'='*50}")
+    print("Napari Labeling Tool")
+    print(f"{'='*50}")
+    print(f"Classes: {args.num_classes}")
+    print("Controls:")
+    print("  B - Paint brush")
+    print("  E - Eraser")
+    print("  1,2,3... - Select label class")
+    print("  Scroll - Navigate slices")
+    print(f"\nClose Napari to save masks to: {args.output}")
+    print(f"{'='*50}\n")
+
+    napari.run()
+
+    # Save on close
+    final_masks = labels_layer.data
+    offset = start if start else 0
+    save_masks(final_masks, args.output, offset=offset)
+
+
+if __name__ == "__main__":
+    main()