segval.py

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Validate a trained YOLOv5 model accuracy on a custom dataset

Usage:
    $ python path/to/val.py --data coco128.yaml --weights yolov5s.pt --img 640
"""

import argparse
import json
import os
import sys
from pathlib import Path
from threading import Thread
from utils.general import LOGGER

import numpy as np
import torch
from tqdm import tqdm



class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.initialized = False
        self.val = None
        self.avg = None
        self.sum = None
        self.count = None

    def initialize(self, val, weight):
        self.val = val
        self.avg = val
        self.sum = val * weight
        self.count = weight
        self.initialized = True

    def update(self, val, weight=1):
        if not self.initialized:
            self.initialize(val, weight)
        else:
            self.add(val, weight)

    def add(self, val, weight):
        self.val = val
        self.sum += val * weight
        self.count += weight
        self.avg = self.sum / self.count

    def value(self):
        return self.val

    def average(self):
        return self.avg


def get_confusion_matrix(label, pred, size, num_class, ignore=-1):
    """
    Calcute the confusion matrix by given label and pred
    """
    output = pred.cpu().numpy().transpose(0, 2, 3, 1)
    seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
    seg_gt = np.asarray(
        label.cpu().numpy()[:, :size[-2], :size[-1]], dtype=np.int)

    ignore_index = seg_gt != ignore
    seg_gt = seg_gt[ignore_index]
    seg_pred = seg_pred[ignore_index]

    index = (seg_gt * num_class + seg_pred).astype('int32')
    label_count = np.bincount(index)
    confusion_matrix = np.zeros((num_class, num_class))

    for i_label in range(num_class):
        for i_pred in range(num_class):
            cur_index = i_label * num_class + i_pred
            if cur_index < len(label_count):
                confusion_matrix[i_label,
                                 i_pred] = label_count[cur_index]
    return confusion_matrix


def validate(testloader, model, numcls, criterion):
    model.eval()
    ave_loss = AverageMeter()
    confusion_matrix = np.zeros((numcls, numcls))

    s = ('%11s' * 3) % ('Images', 'Ave Loss', 'mIOU')
    pbar = tqdm(testloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')  # progress bar

    with torch.no_grad():
        for idx, batch in enumerate(pbar):
            image, label, _, _ = batch
            size = label.size()
            image = image.cuda().float()/255
            label = label.long().cuda()

            # 注意这里要修改一下，主要是因为在Model这个类当中，返回值变成了两个，所以这里要添加上isdomain并且除此之外还要给返回值一个位置
            # 不然，容易报错
            out, _ = model(image, isdomain=False)  # hanfujun
            pred = out[1]
            losses = criterion(pred, label)

            confusion_matrix += get_confusion_matrix(label, pred, size, numcls, 255)

            loss = losses.mean()

            reduced_loss = loss
            ave_loss.update(reduced_loss.item())

        confusion_matrix = torch.from_numpy(confusion_matrix).cuda()
        reduced_confusion_matrix = confusion_matrix
        confusion_matrix = reduced_confusion_matrix.cpu().numpy()

        pos = confusion_matrix.sum(1)
        res = confusion_matrix.sum(0)
        tp = np.diag(confusion_matrix)
        IoU_array = (tp / np.maximum(1.0, pos + res - tp))
        mean_IoU = IoU_array.mean()
    LOGGER.info('%11g' * 3 % (len(testloader), ave_loss.average(), mean_IoU))

    return ave_loss.average(), mean_IoU, IoU_array