eval.py

"""
 * @Author: YBIO
"""
from tqdm import tqdm
import network
import utils
import os
import random
import argparse
import numpy as np
from torch.utils import data
from datasets import VOCSegmentation, ISPRSSegmentation
from utils import ext_transforms as et
from metrics import StreamSegMetrics
import torch
import torch.nn as nn
from utils.tasks import get_tasks
from PIL import Image
import pandas as pd

torch.backends.cudnn.benchmark = True


def get_argparser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_root", type=str, default='/path/to/data',
                        help="path to Dataset")
    parser.add_argument("--dataset", type=str, default='ISPRS', help='Name of dataset')
    parser.add_argument("--num_classes", type=int, default=None, help="num classes (default: None)")
    parser.add_argument("--model", type=str, default='deeplabv3_resnet101',
                        help='model name')
    parser.add_argument("--separable_conv", action='store_true', default=False,
                        help="apply separable conv to decoder and aspp")
    parser.add_argument("--output_stride", type=int, default=16)
    parser.add_argument("--amp", action='store_true', default=False)
    parser.add_argument("--freeze", action='store_true', default=False)
    parser.add_argument("--test_only", action='store_true', default=False)
    parser.add_argument("--total_itrs", type=int, default=30e3,
                        help="epoch number (default: 30k)")
    parser.add_argument("--train_epoch", type=int, default=0,
                        help="epoch number (default: 0")
    parser.add_argument("--curr_itrs", type=int, default=0)
    parser.add_argument("--lr", type=float, default=0.01,
                        help="learning rate (default: 0.01)")
    parser.add_argument("--lr_policy", type=str, default='warm_poly',
                        help="learning rate scheduler policy")
    parser.add_argument("--step_size", type=int, default=10000)
    parser.add_argument("--crop_val", action='store_true', default=False,
                        help='crop validation (default: False)')
    parser.add_argument("--batch_size", type=int, default=16,
                        help='batch size (default: 16)')
    parser.add_argument("--val_batch_size", type=int, default=4,
                        help='batch size for validation (default: 4)')
    parser.add_argument("--crop_size", type=int, default=513)
    parser.add_argument("--ckpt", default=None, type=str,
                        help="restore from checkpoint")
    parser.add_argument("--loss_type", type=str, default='bce_loss',
                        help="loss type (default: False)")
    parser.add_argument("--KD_loss_type", type=str, default='KLDiv_loss',
                        help="KD loss type for ret features")
    parser.add_argument("--use_KD_layer_weight", action='store_true', default=False,
                        help='Whether to apply layer weight for ret feature distillation (default: False)')
    parser.add_argument("--use_KD_class_weight", action='store_true', default=False,
                        help='Whether to apply class weight for ret feature distillation (default: False)')
    parser.add_argument("--gpu_id", type=str, default='0',
                        help="GPU ID")
    parser.add_argument("--weight_decay", type=float, default=1e-5,
                        help='weight decay (default: 1e-5)')
    parser.add_argument("--random_seed", type=int, default=1,
                        help="random seed (default: 1)")
    parser.add_argument("--print_interval", type=int, default=10,
                        help="print interval of loss (default: 10)")
    parser.add_argument("--val_interval", type=int, default=100,
                        help="epoch interval for eval (default: 100)")
    parser.add_argument("--download", action='store_true', default=False,
                        help="download datasets")
    parser.add_argument("--pseudo", action='store_true', default=False)
    parser.add_argument("--pseudo_thresh", type=float, default=0.7)
    parser.add_argument("--task", type=str, default='15-1')
    parser.add_argument("--curr_step", type=int, default=0)
    parser.add_argument("--overlap", action='store_true', default=False)
    parser.add_argument("--mem_size", type=int, default=0)
    parser.add_argument("--bn_freeze", action='store_true', default=False)
    parser.add_argument("--w_transfer", action='store_true', default=False)
    parser.add_argument("--unknown", action='store_true', default=False)
    parser.add_argument("--generate", action='store_true', default=False)
    parser.add_argument("--vision", action='store_true', default=False)
    parser.add_argument("--ddpm_root", type=str, default=None)
    parser.add_argument("--name", type=str, default='None', help="name of the run")

    return parser


def get_dataset(opts):
    if opts.crop_val:
        val_transform = et.ExtCompose([
            et.ExtResize(opts.crop_size),
            et.ExtCenterCrop(opts.crop_size),
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])
    else:
        val_transform = et.ExtCompose([
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])
    if opts.dataset == 'ISPRS':
        dataset = ISPRSSegmentation
    elif opts.dataset == 'dresden':
        dataset = VOCSegmentation
    else:
        raise NotImplementedError
    dataset_dict = {}
    dataset_dict['test'] = dataset(opts=opts, image_set='test', transform=val_transform, cil_step=opts.curr_step)
    return dataset_dict


def validate(opts, model, loader, device, metrics):
    metrics.reset()
    ret_samples = None  # ["intestinal_veins07image49", "ureter04image22", "stomach03image10", "ureter04image19"]
    color_map = {
        0: (0, 0, 0),  # BG
        1: (255, 0, 0),  # red: abdominal wall
        2: (0, 255, 0),  # colon
        3: (0, 0, 255),  # liver
        4: (255, 255, 0),  # pancreas
        5: (255, 0, 255),  # small intestine
        6: (0, 255, 255),  # spleen
        7: (128, 128, 0),  # stomach
        8: (0, 128, 128),  # ureter
        9: (128, 0, 128),  # vesicular glands
        10: (96, 96, 96),  # inferior mesenteric artery
        11: (192, 156, 156),  # intestinal veins
        12: (192, 0, 0),  # 255
        13: (256, 0, 0),  # 255
        255: (0, 0, 0)  # 255
    }
    ious = {}
    for i in range(15):
        ious[i] = [0, 0, 0, 0]
    os.makedirs(f'./eval/{opts.name}', exist_ok=True)
    results = {"name": [], "gt": [], "IoU": []}
    with torch.no_grad():
        for i, (images, labels, _, name) in enumerate(tqdm(loader)):
            flag = 0
            for j in range(len(images)):
                save_name = name[0][j].split('/')[-3] + name[0][j].split('/')[-2] + name[0][j].split('/')[-1]
                save_name = save_name[:-4]
                if ret_samples is not None and save_name in ret_samples:
                    print(save_name)
                    flag = 1
            if ret_samples is not None and flag == 0:
                continue
            images = images.to(device, dtype=torch.float32, non_blocking=True)
            labels = labels.to(device, dtype=torch.long, non_blocking=True)
            ret_features, outputs = model(images)

            output_iou = True
            if output_iou:
                outputs_pred = outputs.detach()
                outputs_pred = torch.argmax(outputs_pred, dim=1)
                outputs_pred = outputs_pred.cpu().numpy()
                labels_imgs = labels.cpu().numpy()

                for j in range(outputs_pred.shape[0]):
                    save_name = name[0][j].split('/')[-3] + name[0][j].split('/')[-2] + name[0][j].split('/')[-1]
                    save_name = save_name[:-4]
                    cls = name[1][j].item()
                    output = outputs_pred[j]
                    labels_img = labels_imgs[j]

                    labels_img[labels_img == 255] = 0
                    # IoU
                    foreground1 = (output == cls).astype(int)
                    foreground2 = (labels_img == cls).astype(int)
                    intersection = np.logical_and(foreground1, foreground2).sum()
                    union = np.logical_or(foreground1, foreground2).sum()
                    iou = intersection / union
                    ious[cls][0] += intersection
                    ious[cls][1] += union
                    ious[cls][2] += iou
                    ious[cls][3] += 1
                    results["name"].append(name[0][j])
                    results["gt"].append(cls)
                    results["IoU"].append(iou)

                    if opts.vision:
                        color_image = np.zeros(output.shape + (3,))
                        color_target = np.zeros(labels_img.shape + (3,))
                        for key, color in color_map.items():
                            color_image[output == key] = color
                            color_target[labels_img == key] = color

                        image = images[j].cpu().numpy().transpose(1, 2, 0)
                        image = (image - image.min()) / (image.max() - image.min()) * 255

                        show_image = np.concatenate([image, color_image, color_target], axis=1)
                        print(save_name, np.unique(output), np.unique(labels_img), output.shape, labels_img.shape, cls, iou)
                        pred_img = Image.fromarray(show_image.astype(np.uint8))
                        pred_img.save(f'./eval/{opts.name}/{iou:.4f}_{save_name}_pred.png')
                        print(f'./eval/{opts.name}/{iou:.4f}_{save_name}_pred.png')

            if opts.loss_type == 'bce_loss':
                outputs = torch.sigmoid(outputs)
            elif opts.loss_type == 'focal_loss':
                outputs = torch.sigmoid(outputs)
            else:
                outputs = torch.softmax(outputs, dim=1)
            preds = outputs.detach().max(dim=1)[1].cpu().numpy()
            targets = labels.cpu().numpy()
            for j in range(preds.shape[0]):
                cls = name[1][j].item()
                preds[j] = np.where(preds[j] == cls, cls, 0)
                targets[j] = np.where(targets[j] == cls, cls, 0)
            metrics.update(targets, preds)
        score = metrics.get_results()

        # output the mIoU for each class and write to a txt file
        all_Gious = []
        all_Mious = []
        ious_text = ""
        for key, value in ious.items():
            if value[1] == 0:
                continue
            print(f'class {key} : mIoU : %.6f' % (value[0] / value[1]), (value[2] / value[3]), value[0], value[1])
            ious_text += f'class {key} : mIoU : %.6f' % (value[0] / value[1]) + '\n'
            all_Gious.append(value[0] / value[1])
            all_Mious.append(value[2] / value[3])
        print(f'GmIoU : %.6f' % np.mean(all_Gious))
        ious_text += f'GmIoU : %.6f' % np.mean(all_Gious) + '\n'
        print(f'mIoU : %.6f' % np.mean(all_Mious))
        ious_text += f'mIoU : %.6f' % np.mean(all_Mious) + '\n'

        # write to a csv file
        df = pd.DataFrame(results)
        df.to_csv(f'./eval/{opts.name}/results.csv', index=False)

    return score, ious_text


def main(opts):
    os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
    target_cls = get_tasks(opts.dataset, opts.task, opts.curr_step)
    opts.num_classes = [len(get_tasks(opts.dataset, opts.task, step)) for step in range(opts.curr_step + 1)]
    if opts.unknown:
        opts.num_classes = [1, 1, opts.num_classes[0] - 1] + opts.num_classes[1:]
    fg_idx = 1 if opts.unknown else 0
    curr_idx = [
        sum(len(get_tasks(opts.dataset, opts.task, step)) for step in range(opts.curr_step)),
        sum(len(get_tasks(opts.dataset, opts.task, step)) for step in range(opts.curr_step + 1))
    ]
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    torch.manual_seed(opts.random_seed)
    np.random.seed(opts.random_seed)
    random.seed(opts.random_seed)
    if 'vit' in opts.model:
        model = network.Segmenter(backbone=opts.model, num_classes=opts.num_classes, pretrained=True)
    else:
        model_map = {
            'deeplabv3_resnet101': network.deeplabv3_resnet101,
        }
        model = model_map[opts.model](num_classes=opts.num_classes, output_stride=opts.output_stride,
                                      bn_freeze=opts.bn_freeze)
        if opts.separable_conv and 'plus' in opts.model:
            network.convert_to_separable_conv(model.classifier)
        utils.set_bn_momentum(model.backbone, momentum=0.01)
    metrics = StreamSegMetrics(sum(opts.num_classes) - 1 if opts.unknown else sum(opts.num_classes),
                               dataset=opts.dataset)
    if opts.overlap:
        ckpt_str = "checkpoints/%s_%s_%s_step_%d_overlap.pth"
    else:
        ckpt_str = "checkpoints/%s_%s_%s_step_%d_disjoint.pth"
    model = nn.DataParallel(model)
    model = model.to(device)
    dataset_dict = get_dataset(opts)
    test_loader = data.DataLoader(
        dataset_dict['test'], batch_size=opts.val_batch_size, shuffle=False, num_workers=4, pin_memory=True)
    report_dict = dict()
    best_ckpt = ckpt_str % (opts.model, opts.dataset, opts.task, opts.curr_step)
    best_ckpt = opts.ckpt if opts.ckpt else best_ckpt
    checkpoint = torch.load(best_ckpt, map_location=torch.device('cpu'))
    model.module.load_state_dict(checkpoint["model_state"], strict=False)
    model.eval()
    test_score, results = validate(opts=opts, model=model, loader=test_loader,
                          device=device, metrics=metrics)
    print(metrics.to_str(test_score))
    with open(f'./eval/{opts.name}/results.txt', 'w') as f:
        f.write(results)
        f.write(metrics.to_str(test_score))
    report_dict[f'a_miou'] = test_score['Mean IoU']
    class_iou = list(test_score['Class IoU'].values())
    class_dice = list(test_score['Class Dice'].values())
    first_cls = len(get_tasks(opts.dataset, opts.task, 0))
    report_dict[f'b_mIoU'] = np.mean(class_iou[:first_cls])
    report_dict[f'n_mIoU'] = np.mean(class_iou[first_cls:])
    print(f" 1 - {first_cls - 1} : mIoU : %.6f" % np.mean(class_iou[1:first_cls]))
    print(f" 1 - {first_cls - 1} : mDice : %.6f" % np.mean(class_dice[1:first_cls]))
    print(f" {first_cls} - {len(class_iou) - 1}: mIoU : %.6f" % np.mean(class_iou[first_cls:]))
    print(f" {first_cls} - {len(class_iou) - 1}: mDice : %.6f" % np.mean(class_dice[first_cls:]))
    print(f" all_miou : %.6f" % np.mean(class_iou[1:]))
    print(f" all_mDice : %.6f" % np.mean(class_dice[1:]))


if __name__ == '__main__':
    opts = get_argparser().parse_args()
    total_step = len(get_tasks(opts.dataset, opts.task))
    opts.curr_step = total_step - 1
    main(opts)