| import torch |
| import numpy as np |
| from torch.utils.data import DataLoader |
| from tqdm import tqdm |
| import yaml |
| import wandb |
| import cv2 |
| from monai.losses import DiceCELoss |
| from monai.metrics import DiceMetric |
| from monai.transforms import Compose, Activations, AsDiscrete |
| from monai.networks.nets import UNet, SwinUNETR |
| from dataloading.dataloader2D import NiftiSegmentationDataset |
| import os |
|
|
| |
| |
| |
| with open("/workspace/Segmentation/config2d.yaml", "r") as f: |
| config = yaml.safe_load(f) |
|
|
| output_dir = config['output_dir'] |
|
|
|
|
| device = torch.device(config['device'] if torch.cuda.is_available() else "cpu") |
|
|
| |
| |
| |
|
|
| model =SwinUNETR(img_size=(256, 256), in_channels=3, out_channels=config['num_classes'], spatial_dims=2).to(device) |
|
|
| model.load_state_dict(torch.load("/workspace/Segmentation/checkpoints/best_model.pth", map_location=device)) |
| model.eval() |
|
|
| |
| |
| |
| test_dataset = NiftiSegmentationDataset(config['csv_file_train'], channel_keys=config['channel_keys'], augment=False) |
| test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) |
|
|
| |
| |
| |
| dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False) |
| post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) |
|
|
| |
| |
| |
| with torch.no_grad(): |
| for i, (x, y) in enumerate(tqdm(test_loader, desc="Testing")): |
| x, y = x.to(device), y.to(device) |
|
|
| logits = model(x) |
| probs = post_pred(logits) |
| dice_metric(y_pred=probs, y=y) |
|
|
| img_np = x[0].detach().cpu().numpy() |
| gt_np = y[0].detach().cpu().numpy().astype(np.uint8) |
| pred_np = torch.argmax(probs[0], dim=0).detach().cpu().numpy().astype(np.uint8) |
|
|
| if img_np.shape[0] == 1: |
| img_np = img_np[0] |
|
|
| gt_np = np.squeeze(gt_np).astype(np.uint8) |
| pred_np = np.squeeze(pred_np).astype(np.uint8) |
|
|
| img_norm = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8) |
| img_uint8 = (img_norm * 255).astype(np.uint8) |
| img_rgb = cv2.cvtColor(img_uint8, cv2.COLOR_GRAY2BGR) |
|
|
| contour_img = img_rgb.copy() |
| contours_gt, _ = cv2.findContours(gt_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| cv2.drawContours(contour_img, contours_gt, -1, (0, 255, 0), 2) |
|
|
| contours_pred, _ = cv2.findContours(pred_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| cv2.drawContours(contour_img, contours_pred, -1, (0, 0, 255), 2) |
|
|
| |
| filename = os.path.basename(test_dataset.samples[i]["image_path"]) |
| save_path = os.path.join(output_dir, f"overlay_{i}_{filename.replace('.nii.gz', '.png')}") |
| cv2.imwrite(save_path, contour_img) |
|
|
| print(dice_metric.aggregate().item()) |
