inference.py

import argparse
import json
import numpy as np
import os
import pandas as pd
import torch
from glob import glob
from scipy.ndimage.morphology import binary_fill_holes
from skimage.io import imread
from skimage.morphology import disk, binary_erosion, label
from skimage.transform import downscale_local_mean
from tqdm import tqdm

from dataset import TomoDetectionDataset as Dataset
from dense_yolo import DenseYOLO
from subsets import data_frame_subset

cell_size = Dataset.cell_size
# larger grid size for inference to run inference on full image without cropping
img_height = cell_size * 12
img_width = cell_size * 9
grid_size = (img_height // cell_size, img_width // cell_size)
anchor = Dataset.anchor


def main(args, config):
    data_frame = data_frame_subset(
        args.data_views, args.data_boxes, args.subset, seed=args.seed
    )
    pred_data_frame = pd.DataFrame()

    if args.only_biopsied:
        data_frame = data_frame[(data_frame["Benign"] == 1) | (data_frame["Cancer"] == 1)]

    with torch.set_grad_enabled(False):
        yolo = DenseYOLO(img_channels=1, out_channels=Dataset.out_channels, **config)

        if args.multi_gpu and torch.cuda.device_count() > 1:
            device = torch.device("cuda:0")
            yolo = torch.nn.DataParallel(yolo)
        else:
            device = torch.device("cpu" if not torch.cuda.is_available() else args.device)

        yolo.to(device)

        state_dict = torch.load(args.weights)
        yolo.load_state_dict(state_dict)
        yolo.eval()
        yolo.to(device)

        for index, row in tqdm(data_frame.iterrows(), total=len(data_frame)):
            pid = str(row["PatientID"]).zfill(5)
            sid = row["StudyUID"]
            view = str(row["View"])

            view_template = "{}TomosynthesisReconstruction_*_.png".format(view.upper())
            view_files = glob(os.path.join(args.images, pid, sid, view_template))

            batch = []
            volume = []
            pred_view = np.zeros((len(view_files), 5) + grid_size)

            for slice_n in range(len(view_files)):
                batch.append(
                    read_slice_image(
                        pid, sid, view, slice_n, args.images, args.downscale
                    )
                )
                volume.append(batch[-1][0])

                if len(batch) >= args.batch_size:
                    y_pred = predict(yolo, batch, device)
                    pred_view[slice_n + 1 - len(batch) : slice_n + 1] = y_pred
                    batch = []

            if len(batch) > 0:
                y_pred = predict(yolo, batch, device)
                pred_view[-len(batch) :] = y_pred

            pred_view = average_predictions(pred_view, view_split=args.view_split)
            if args.keep_splits > 0:
                pred_view = filter_by_score(pred_view, keep=args.keep_splits)

            slice_span = len(volume) / args.view_split
            df_view_bboxes = pred2bboxes(
                pred_view, slice_span=slice_span, threshold=args.pred_threshold
            )

            df_view_bboxes = remove_empty_boxes(df_view_bboxes, np.array(volume))

            df_view_bboxes["PatientID"] = pid
            df_view_bboxes["StudyUID"] = sid
            df_view_bboxes["View"] = view
            pred_data_frame = pred_data_frame.append(
                df_view_bboxes, ignore_index=True, sort=False
            )

    # rescale boxes to original images size
    pred_data_frame["X"] = pred_data_frame["X"] * args.downscale
    pred_data_frame["Y"] = pred_data_frame["Y"] * args.downscale
    pred_data_frame["Width"] = pred_data_frame["Width"] * args.downscale
    pred_data_frame["Height"] = pred_data_frame["Height"] * args.downscale

    pred_data_frame = pred_data_frame[
        ["PatientID", "StudyUID", "View", "Score", "Z", "X", "Y", "Depth", "Width", "Height"]
    ]
    pred_data_frame[["X", "Y", "Z", "Width", "Height", "Depth"]] = pred_data_frame[
        ["X", "Y", "Z", "Width", "Height", "Depth"]
    ].astype(int)

    pred_data_frame.to_csv(args.predictions, index=False)


def predict(model, batch, device):
    batch_tensor = torch.from_numpy(np.array(batch))
    batch_tensor = batch_tensor.to(device)
    y_pred_device = model(batch_tensor)
    y_pred = y_pred_device.cpu().numpy()
    return np.squeeze(y_pred)


def read_slice_image(pid, sid, view, slice_n, images_dir, downscale):
    filename = "{}TomosynthesisReconstruction_{}_.png".format(view.upper(), slice_n)
    image_path = os.path.join(images_dir, pid, sid, filename)
    img = _imread(image_path, downscale=downscale, flip="R" in view.upper())

    if img.shape[0] < img_height:
        pad_y = img_height - img.shape[0]
        img = np.pad(img, ((0, pad_y), (0, 0)), mode="constant")
    elif img.shape[0] > img_height:
        img = img[:img_height, :]

    if img.shape[1] < img_width:
        pad_x = img_width - img.shape[1]
        img = np.pad(img, ((0, 0), (0, pad_x)), mode="constant")
    elif img.shape[1] > img_width:
        img = img[:, :img_width]

    # normalize
    img = img.astype(np.float32) / np.max(img)

    # fix dimensions (N, C, H, W)
    img = img[..., np.newaxis]
    img = img.transpose((2, 0, 1))

    return img


def _imread(imgpath, downscale, flip=False):
    image = imread(imgpath)
    if downscale != 1:
        image = downscale_local_mean(image, (downscale, downscale))
    if flip:
        image = np.fliplr(image).copy()
    image = _preprocess(image)
    return image


def _preprocess(image, erosion=5):
    mask = _mask(image, erosion=erosion)
    image = mask * image
    return image


def _mask(image, erosion=10):
    mask = image > 0
    mask = np.pad(mask, ((0, 0), (1, 0)), mode="constant", constant_values=1)
    mask = binary_fill_holes(mask)
    mask = mask[:, 1:]
    mask = binary_erosion(mask, disk(erosion))
    cc = label(mask, background=0)
    lcc = np.argmax(np.bincount(cc.flat)[1:]) + 1
    mask = cc == lcc
    return mask


def _mean_filter(image, filter_size=4):
    fs = filter_size
    yy, xx = np.nonzero(image >= np.max(image) * 0.99)
    image_out = image
    for y, x in zip(yy, xx):
        neighborhood = image[max(0, y - fs) : y + fs, max(0, x - fs) : x + fs]
        image_out[y, x] = np.mean(neighborhood)
    return image_out


def average_predictions(pred_view, view_split=4):
    pred_view_avg = np.zeros((view_split, 5) + grid_size)
    slice_span = int(pred_view.shape[0] / view_split)
    for i in range(view_split):
        pred_view_avg[i] = np.mean(
            pred_view[i * slice_span : (i + 1) * slice_span], axis=0
        )
    return pred_view_avg


def filter_by_score(pred_view, keep):
    if keep >= pred_view.shape[0]:
        return pred_view
    for i in range(pred_view.shape[-2]):
        for j in range(pred_view.shape[-1]):
            pred_cell = pred_view[:, 0, i, j]
            threshold = sorted(pred_cell.flat, reverse=True)[keep]
            for k in range(pred_view.shape[0]):
                if pred_view[k, 0, i, j] <= threshold:
                    pred_view[k, 0, i, j] = 0.0
    return pred_view


def pred2bboxes(pred, slice_span, threshold=None):
    # box: upper-left corner + width + height + first slice + depth
    np.nan_to_num(pred, copy=False)
    obj_th = pred[:, 0, ...]
    if threshold is None:
        threshold = min(0.0001, np.max(obj_th) * 0.5)
    obj_th[obj_th < threshold] = 0
    z, y, x = np.nonzero(obj_th)
    scores = []
    xs = []
    ys = []
    hs = []
    ws = []
    for i in range(len(z)):
        scores.append(pred[z[i], 0, y[i], x[i]])
        h = int(anchor[0] * pred[z[i], 3, y[i], x[i]] ** 2)
        hs.append(h)
        w = int(anchor[0] * pred[z[i], 4, y[i], x[i]] ** 2)
        ws.append(w)
        y_offset = pred[z[i], 1, y[i], x[i]]
        y_mid = y[i] * cell_size + (cell_size / 2) + (cell_size / 2) * y_offset
        ys.append(int(y_mid - h / 2))
        x_offset = pred[z[i], 2, y[i], x[i]]
        x_mid = x[i] * cell_size + (cell_size / 2) + (cell_size / 2) * x_offset
        xs.append(int(x_mid - w / 2))

    zs = [s * slice_span for s in z]
    df_dict = {
        "Z": zs,
        "X": xs,
        "Y": ys,
        "Width": ws,
        "Height": hs,
        "Depth": [slice_span] * len(zs),
        "Score": scores,
    }
    df_bboxes = pd.DataFrame(df_dict)
    df_bboxes.sort_values(by="Score", ascending=False, inplace=True)
    return df_bboxes


def remove_empty_boxes(df, volume):
    # box: upper-left corner + width + height + first slice + depth
    empty_indices = []
    for index, box in df.iterrows():
        w = int(box["Width"])
        h = int(box["Height"])
        d = int(box["Depth"])
        x = int(max(box["X"], 0))
        y = int(max(box["Y"], 0))
        z = int(max(box["Z"], 0))
        box_volume = volume[z : z + d, y : y + h, x : x + w]
        if np.sum(box_volume == 0) > 0.5 * w * h * d:
            empty_indices.append(index)
    df = df.drop(index=empty_indices)
    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Running inference using trained YOLO model for cancer detection in Duke DBT volumes"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="device for testing (default: cuda:1)",
    )
    parser.add_argument(
        "--multi-gpu",
        default=False,
        action="store_true",
        help="flag to train on multiple gpus",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=16,
        help="input batch size for testing (default: 16)",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default="config_default.json",
        help="network config file (see: config_default.json)",
    )
    parser.add_argument(
        "--weights", type=str, required=True, help="file with saved weights"
    )
    parser.add_argument(
        "--data-views",
        type=str,
        default="/data/data_train_v2.csv",
        help="csv file listing training/test views together with category label",
    )
    parser.add_argument(
        "--data-boxes",
        type=str,
        default="/data/bboxes_v2.csv",
        help="csv file defining ground truth bounding boxes",
    )
    parser.add_argument(
        "--images",
        type=str,
        default="/data/TomoImages/",
        help="root folder with images",
    )
    parser.add_argument(
        "--predictions",
        type=str,
        required=True,
        help="output file path with predictions",
    )
    parser.add_argument(
        "--view-split",
        type=int,
        default=2,
        help="number of view parts for averaging predictions (default: 2)",
    )
    parser.add_argument(
        "--keep-splits",
        type=int,
        default=0,
        help="number of averaged view splits to keep after filtering (default: 0=view-split)",
    )
    parser.add_argument(
        "--pred-threshold",
        type=float,
        default=0.0001,
        help="threshold for minimum box prediction confidence (default: 0.0001)",
    )
    parser.add_argument(
        "--only-biopsied",
        default=False,
        action="store_true",
        help="flag to use only biopsied cases",
    )
    parser.add_argument(
        "--subset",
        type=str,
        default="validation",
        help="subset to run inference on [all|train|validation|test] (default: validation)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="random seed for validation split (default: 42)",
    )
    parser.add_argument(
        "--downscale",
        type=int,
        default=2,
        help="input image downscale factor used to upscale boxes to original scale (default 2)",
    )
    args = parser.parse_args()

    with open(args.config, "r") as fp:
        config = json.load(fp)

    main(args, config)