init

lib/evaluation/segm_eval_base.py

@@ -0,0 +1,208 @@
+import os
+import argparse
+
+import numpy as np
+import pandas as pd
+from skimage.transform import resize
+
+from skimage import io
+from multiprocessing import Pool
+from functools import partial
+import logging
+import datetime
+
+def get_confusion_matrix(pred, gt, num_class):
+    assert pred.shape == gt.shape,  f"pred.shape: {pred.shape} != gt.shape: {gt.shape}"
+    mask = (gt >= 0) & (gt < num_class) # 去掉为0的背景类别
+    label = num_class * gt[mask] + pred[mask]
+    count = np.bincount(label, minlength=num_class**2)
+    confusion_matrix = count.reshape(num_class, num_class)
+    return confusion_matrix
+
+def get_miou(confusion_matrix):
+    diagonal_elements = np.diag(confusion_matrix)
+    column_sums = np.sum(confusion_matrix, axis=0)
+    row_sums = np.sum(confusion_matrix, axis=1)
+    ious = diagonal_elements/(column_sums + row_sums - diagonal_elements)
+    m_iou = np.nanmean(ious)
+    return m_iou
+
+def get_mprecison(confusion_matrix):
+    diagonal_elements = np.diag(confusion_matrix)
+    column_sums = np.sum(confusion_matrix, axis=0)
+    precisions = diagonal_elements / (column_sums + 1e-06) 
+    m_precision = np.nanmean(precisions)
+    return m_precision
+
+def get_mrecall(confusion_matrix):
+    diagonal_elements = np.diag(confusion_matrix)
+    row_sums = np.sum(confusion_matrix, axis=1)
+    recalls= diagonal_elements / (row_sums + 1e-06)
+    m_recall = np.nanmean(recalls)
+    return m_recall 
+
+def get_macc(confusion_matrix):
+    '''
+    acc = tp/tp+fn 就是recall
+    '''
+    m_recall = get_mrecall(confusion_matrix)
+    return m_recall 
+
+def get_per_class_iou(confusion_matrix):
+    intersection = np.diag(confusion_matrix)
+    union = np.sum(confusion_matrix, axis=0) + np.sum(confusion_matrix, axis=1) - intersection
+    iou = intersection / (union.astype(np.float32) + 1e-6)
+    return iou
+
+def get_per_class_acc(confusion_matrix):
+    total_acc = np.diag(confusion_matrix) / (np.sum(confusion_matrix, axis=1).astype(np.float32) + 1e-6)
+    return total_acc
+
+
+def post_process_segm_output(segm, colors, dist_type='abs'):
+    """
+    Post-processing to turn output segm image to class index map using NumPy
+    Args:
+        segm: (H, W, 3)
+    Returns:
+        class_map: (H, W)
+    """
+    palette = np.array(colors)
+    segm = segm.astype(np.float32)  # (h, w, 3)
+    h, w, k = segm.shape[0], segm.shape[1], palette.shape[0]
+    if dist_type == 'abs':
+        dist = np.abs(segm.reshape(h, w, 1, 3) - palette.reshape(1, 1, k, 3))  # (h, w, k)
+    elif dist_type == 'square':
+        dist = np.power(segm.reshape(h, w, 1, 3) - palette.reshape(1, 1, k, 3), 2)  # (h, w, k)
+    elif dist_type == 'mean':
+        dist_abs = np.abs(segm.reshape(h, w, 1, 3) - palette.reshape(1, 1, k, 3))  # (h, w, k)
+        dist_square = np.power(segm.reshape(h, w, 1, 3) - palette.reshape(1, 1, k, 3), 2)  # (h, w, k)
+        dist = (dist_abs + dist_square) / 2.
+    else:
+        raise NotImplementedError
+    
+    dist = np.sum(dist, axis=-1)
+    pred = np.argmin(dist, axis=-1).astype(np.int)
+    return pred
+    
+def get_args_parser():
+    parser = argparse.ArgumentParser('semantic segmentation evaluation', add_help=False)
+    parser.add_argument('--pred_dir', type=str, help='dir to pred', required=True)
+    parser.add_argument('--gt_dir', type=str, help='dir to gt', required=True)
+    parser.add_argument('--gt_list_path', type=str, help='dir to gt_list_path', required=True)
+    parser.add_argument('--gt_suffix', type=str, help='suffix to gt', required=True)
+    parser.add_argument('--dataset_name', type=str, help='dataset name', required=True)
+    parser.add_argument('--model_name', type=str, help='model name', required=True)
+    parser.add_argument('--dist_type', type=str, help='dist type',
+                        default='abs', choices=['abs', 'square', 'mean'])
+    return parser.parse_args()
+
+def process_file(file_dict, pred_dir, gt_dir, args, num_class):
+    filename = file_dict['file_name']
+    file_cls = file_dict['file_cls']
+
+    gt = io.imread(os.path.join(gt_dir, filename))
+    gt_index = gt.copy()
+    gt_index[gt_index != file_cls] = 0
+    gt_index[gt_index == file_cls] = 1
+
+    try:
+        pred = io.imread(os.path.join(pred_dir, filename.replace('.png', f'-{file_cls}.png')))
+        pred = resize(pred, gt.shape[-2:], anti_aliasing=False, mode='reflect', order=0)
+        
+        if len(pred.shape) == 3:
+            pred_index = pred[:,:,0].copy()
+        else:
+            pred_index = pred.copy()
+        pred_index[pred_index<=127] = 0
+        pred_index[pred_index>127] = 1
+    except:
+        logging.info(filename.replace('.png', f'_{file_cls}.png'), 'not found!')
+        pred_index = gt_index.copy()
+    
+    pred_index = pred_index.flatten()
+    gt_index = gt_index.flatten()
+    confusion_matrix = get_confusion_matrix(pred_index, gt_index, num_class)
+    return file_cls, confusion_matrix
+
+if __name__ == '__main__':
+    args = get_args_parser()
+    dataset_name = args.dataset_name 
+    pred_dir = args.pred_dir
+    gt_dir = args.gt_dir
+    gt_list_path = args.gt_list_path
+    dist_type = args.dist_type
+    model_name = args.model_name
+
+    current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    os.makedirs('logs/eval', exist_ok=True)
+    logging.basicConfig(
+        level=logging.INFO, 
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler(f'logs/eval/eval_{model_name}_{dataset_name}_{current_time}.log'),
+            logging.StreamHandler() 
+        ]
+    )
+
+    output_folder = os.path.join(pred_dir, f'eval_{dataset_name}')
+    os.makedirs(output_folder, exist_ok=True)
+
+
+    num_class = 2
+
+    with open(gt_list_path, 'r') as f:
+        file_cls_list = f.readlines()
+    
+    file_list = []
+    for i in file_cls_list:
+        i = i.strip()
+        file_name = i[:-3]
+        file_cls = i[-2:]
+        file_list.append({'file_name': file_name, 'file_cls': int(file_cls)})
+    
+    all_pred_labels = []
+    all_gt_labels = []
+
+    process_file_partial = partial(process_file, pred_dir=pred_dir,gt_dir=gt_dir, args=args, num_class=num_class)
+    
+    pool = Pool()
+
+    outputs = pool.map(process_file_partial, file_list)
+    
+    pool.close()
+    pool.join()
+    logging.info(f'len outputs: {len(outputs)}')
+    confusion_matrix_dict = {}
+    for cls, confusion_matrix in outputs:
+        if cls in confusion_matrix_dict.keys():
+            confusion_matrix_dict[cls] += confusion_matrix
+        else:
+            confusion_matrix_dict[cls] = confusion_matrix
+
+    class_list = []
+    iou_list = []
+    acc_list = []
+    for cls, confusion_matrix in confusion_matrix_dict.items():
+        ious = get_per_class_iou(confusion_matrix)
+        accs = get_per_class_acc(confusion_matrix)
+        logging.info(f'cls: {cls}, ious: {ious}, accs: {accs}')
+        class_list.append(cls)
+        iou_list.append(ious[1])
+        acc_list.append(accs[1])
+    
+    miou = np.mean(iou_list)
+    macc = np.mean(acc_list)
+
+    df_metrics = pd.DataFrame({
+        'Class': class_list + ['Mean'],
+        'IoU': iou_list + [miou],
+        'Accuracy':  acc_list + [macc],
+    })
+    pd.set_option('display.float_format', '{:.4f}%'.format)
+    logging.info(df_metrics)
+    pd.reset_option('display.float_format')
+    df_metrics.to_csv(os.path.join(output_folder, 'eval.csv'), index=False, float_format='%.4f')
+    
+