<jupyter_start><jupyter_text>Evaluate pretrained U-net from https://github.com/JoHof/lungmask on Plethora dataset.<jupyter_code>%load_ext autoreload
%autoreload 2
%cd ..
import SimpleITK as sitk
import numpy as np
import matplotlib.pyplot as plt
import skimage

import torch
from torchvision import transforms
from tqdm.notebook import tqdm

from lungmask.mask import get_model
from lungmask.utils import preprocess, postrocessing, reshape_mask

from src.metrics import dice_coeff_vectorized
from src.data.preprocess import ToTensor
from src.data.data_modules import PlethoraDataModule, Covid19DataModule
from src.visualization.plotting import plot_batch
class LungMaskPreprocess:
    def __call__(self, sample):
        img = sample["img"]
        new_img, new_box = preprocess(np.expand_dims(img, axis=0),
                                               resolution=[256,256])
        new_img[new_img > 600] = 600
        new_img = np.divide((new_img + 1024), 1624)
        new_img = np.squeeze(new_img)
        new_box = np.squeeze(new_box)
        return dict(sample, img=new_img, box=new_box, original_img=img)
transform = transforms.Compose([
    LungMaskPreprocess(),
    ToTensor(),
    transforms.Lambda(lambda sample: dict(sample, img=sample["img"].unsqueeze(0))),
])

#dm = PlethoraDataModule(batch_size=16, transform=transform)
dm = Covid19DataModule(batch_size=32, transform=transform)
dm.setup('test')
test_loader = dm.test_dataloader()
device = "cuda:0"
model = get_model('unet', 'R231').to(device)
dice_2d_list = []
dice_3d_list = []
true_buffer = []
pred_buffer = []

res = {}
pbar = tqdm(test_loader)
for batch_idx, batch in enumerate(pbar):
    with torch.no_grad():
        X, y = batch["img"].to(device), batch["mask"].to(device)
        logits = model(X)
        pred_masks = torch.max(logits, 1)[1].detach().cpu().numpy().astype(np.uint8)
    
    boxes = batch["box"].cpu().numpy()
    # post-process masks
    pred_masks = postrocessing(pred_masks)
    pred_masks = np.asarray(
            [reshape_mask(pred_masks[i], batch["box"][i], batch["original_img"].shape[1:]) for i in range(pred_masks.shape[0])],
            dtype=np.uint8)
    pred_masks[pred_masks > 0] = 1
    pred_masks = torch.from_numpy(pred_masks).to(device)
    
    # calculate metrics
    dsc = dice_coeff_vectorized(pred_masks, y, reduce_fn=torch.mean)
    dice_2d_list.append(dsc.item())
    res["dice_2d"] = dsc.item()

    # 3D dice coeff
    slice_idxs = batch["slice_idx"]
    split_idx = torch.where(slice_idxs == 0)[0]
    if len(split_idx) > 1:
        raise RuntimeError(f"there are multiple zeros in slice_idxs: {slice_idxs}")
    split_idx = split_idx.item() if len(split_idx) > 0 else None
    true_buffer.append(y[:split_idx])
    pred_buffer.append(pred_masks[:split_idx])

    # if we have finish iterating over a CT scan, calculate 3D dice
    if (split_idx is not None and batch_idx > 0) or \
            batch_idx == len(test_loader) - 1:
        true_v_mask = torch.cat(true_buffer).reshape(1, -1)
        pred_v_mask = torch.cat(pred_buffer).reshape(1, -1)

        dsc_v = dice_coeff_vectorized(pred_v_mask, true_v_mask, reduce_fn=None)
        dice_3d_list.append(dsc_v.item())
        res["dice_3d"] = dsc_v.item()

        # clear the buffers
        true_buffer = [y[split_idx:]]
        pred_buffer = [pred_masks[split_idx:]]
        del true_v_mask
        del pred_v_mask
    pbar.set_postfix(res)
print(len(dice_3d_list))
print(dice_3d_list)
print(f"2d dice: {np.mean(dice_2d_list):.4f}")
print(f"3d dice: {np.mean(dice_3d_list):.4f}")<jupyter_output>2d dice: 0.9391
3d dice: 0.9822
{'fraction_non_alphanumeric': 0.12340543538546866, 'fraction_numerical': 0.022462562396006656, 'mean_word_length': 3.7523056653491436, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '416549', 'license': 'FTL', 'n_tokens_mistral': 1353, 'n_tokens_neox': 1311, 'n_words': 256, 'path': 'notebooks/7.0-duc-eval-lungmask-repo.ipynb', 'repo_name': 'minhduc0711/lung-segmentation', 'size': '3606'}