Algorithm-Repo/SourceCode/MachineLearningUtils/ObjectDetection/YOLO_V4/compute_test_set_aps.py

import sys
import onnx
import os
import argparse
import numpy as np
import cv2
import onnxruntime
import multiprocessing
from tool.utils import *
from tool.darknet2onnx import *
import glob
onnx_path = 'C:\\Users\\VINNO\\Desktop\\新建文件夹 (2)\\pytorch-YOLOv4-master\\20210824_yolov4_1_224_224_static.onnx'
OUTPUT_PATH = 'C:\\Users\\VINNO\\Desktop\\新建文件夹 (2)\\pytorch-YOLOv4-master\\output\\'


sess_options = onnxruntime.SessionOptions()
#sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL

#控制用于运行模型的线程数 controls the number of threads to use to run the model
#sess_options.intra_op_num_threads = 1
#When sess_options.execution_mode = rt.ExecutionMode.ORT_PARALLEL,
# you can set sess_options.inter_op_num_threads to control the number of threads used to parallelize the execution of the graph (across nodes).
#sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_PARALLEL
#sess_options.inter_op_num_threads = 1
session = onnxruntime.InferenceSession(onnx_path, sess_options)

IN_IMAGE_H = session.get_inputs()[0].shape[2]
IN_IMAGE_W = session.get_inputs()[0].shape[3]
input_name = session.get_inputs()[0].name

# class_id_map = None
class_id_map = [None] * 8
class_id_map[1] = 1
class_id_map[2] = 1
class_id_map[3] = 1
class_id_map[4] = 1
class_id_map[5] = 1
class_id_map[6] = 1
class_id_map[7] = 1

class_list = [None] * 8
class_list[0] = '__background__'
class_list[1] = 'lipomyoma'
class_list[2] = 'BIRADS2'
class_list[3] = 'BIRADS3'
class_list[4] = 'BIRADS4a'
class_list[5] = 'BIRADS4b'
class_list[6] = 'BIRADS4c'
class_list[7] = 'BIRADS5'

#在测试集上计算mAP,绘制ROC曲线等,所有的评估曲线保存在 Output\\TestCurves 文件夹内
def evaluate_detections(c, boxes, gtinfo, dstpath, ovthresh=0.5):
    """
    计算评价指标
    precision=tp/(tp+fp):计算准确率时,以预测框个数为准
    recall=tp/(tp+fn):计算召回率时,以gt框个数为准
    :param c: 当前类在classes中的序号
    :param boxes: 所有的预测结果中属于该类的所有box,
                  类型：list,list长度为测试图像总数,list中每一项为长度为k*5的子list
                  其中,k为当前图像上的预测框个数,每个预测框用'左,上,右,下,置信度'5个值表示
                  若当前图像没有预测到任何框,则子list为空
    :param gtinfo: 所有gt中属于该类的所有box
                   类型：list,list长度为测试图像总数,list中每一项为长度为n*4的子list
                   其中,n为当前图像上gt框的个数,每个gt框用'左,上,右,下'4个值表示
                   若当前图像没有gtbox,则子list为空
    :param dstpath: 输出目录(绘制的roc曲线等保存在指定位置)
    :param ovthresh: 重叠率阈值,当预测框与gtbox的IOU必须大于该值才会被当作是tp
    :return: roc曲线,tp/fp曲线等
    """

    assert len(boxes) == len(gtinfo), "the length of pred boxes and gtinfos should be the same"
    assert ovthresh > 0, "the overlaps between gt boxes and pred boxes should be greater than 0."
    detBboxes = np.zeros((0, 7), dtype=np.float32)
    gtBboxes = np.zeros((0, 6), dtype=np.float32)
    for reading_ind in range(len(boxes)):
        preds = boxes[reading_ind]
        gts = gtinfo[reading_ind]
        assert len(preds) % 5 == 0 and len(gts) % 4 == 0, \
            "incorrect inputs for the pred boxes or gt boxes of image:{} class:{}".format(reading_ind, c)
        num_preds = int(len(preds) / 5)
        num_gts = int(len(gts) / 4)
        if num_preds > 0:
            reshaped_preds = np.reshape(preds, (num_preds, 5))
            imginds = reading_ind * np.ones((num_preds, 1), dtype=np.float32)
            tpsign = np.zeros((num_preds, 1), dtype=np.float32)
            detBboxes = np.vstack((detBboxes, np.hstack((imginds, reshaped_preds, tpsign))))
        if num_gts > 0:
            reshaped_gts = np.reshape(gts, (num_gts, 4))
            imginds = reading_ind * np.ones((num_gts, 1), dtype=np.float32)
            tpsign = np.zeros((num_gts, 1), dtype=np.float32)
            gtBboxes = np.vstack((gtBboxes, np.hstack((imginds, reshaped_gts, tpsign))))

    # 然后计算Tps和Fps
    nd = len(detBboxes)
    ngt = len(gtBboxes)
    if nd > 0:
        # sort by confidence
        sorted_ind = np.argsort(-detBboxes[:, -2])
        detBboxes = detBboxes[sorted_ind, :]
        for d in range(nd):
            detbox = detBboxes[d, :]
            imgind = detbox[0]
            gtind = np.where(gtBboxes[:, 0] == imgind)[0]
            if len(gtind > 0):
                gts = gtBboxes[gtind, :]
                # 计算重叠区域
                ix1 = np.maximum(gts[:, 1], detbox[1])
                iy1 = np.maximum(gts[:, 2], detbox[2])
                ix2 = np.minimum(gts[:, 3], detbox[3])
                iy2 = np.minimum(gts[:, 4], detbox[4])
                iw = np.maximum(ix2 - ix1 + 1., 0.)
                ih = np.maximum(iy2 - iy1 + 1., 0.)
                inters = iw * ih
                # unions
                uni = ((detbox[3] - detbox[1] + 1.) * (detbox[4] - detbox[2] + 1.) +
                       (gts[:, 3] - gts[:, 1] + 1.) * (gts[:, 4] - gts[:, 2] + 1.) - inters)
                overlaps = inters / uni
                ovmax = np.max(overlaps)
                jmax = np.argmax(overlaps)
                if ovmax > ovthresh and ovmax > 0:
                    if gts[jmax, -1] == 0:
                        gtBboxes[gtind[jmax], -1] = 1
                        detBboxes[d, -1] = 1
    # 计算数值
    prec_val = np.sum(detBboxes[:, -1] == 1) / np.maximum(float(nd), np.finfo(np.float32).eps)
    print('prec_val的 tp：{}'.format(np.sum(detBboxes[:, -1] == 1)))
    rec_val = np.sum(gtBboxes[:, -1] == 1) / np.maximum(float(ngt), np.finfo(np.float32).eps)
    print('rec_val的 tp：{}'.format(np.sum(gtBboxes[:, -1] == 1)))
    F1 = 2 * prec_val * rec_val / np.maximum(prec_val + rec_val, np.finfo(np.float32).eps)

    # fp值false positive，实际为负例，预测为正例，为了计算误检
    temp_fp = np.maximum(float(nd), np.finfo(np.float32).eps) - np.sum(detBboxes[:, -1] == 1)

    # 绘制p/r曲线 求mAP
    tp = detBboxes[:, -1]
    fp = 1 - tp
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)
    recalls = tp / np.maximum(float(ngt), np.finfo(np.float32).eps)
    precisions = tp / np.maximum(tp + fp, np.finfo(np.float32).eps)
    mrecalls = np.concatenate(([0.], recalls, [1.]))
    mprecisions = np.concatenate(([0.], precisions, [0.]))
    # compute the precision envelope
    for i in range(mprecisions.size - 1, 0, -1):
        mprecisions[i - 1] = np.maximum(mprecisions[i - 1], mprecisions[i])

    # to calculate area under PR curve, look for points
    # where X axis (recall) changes value
    i = np.where(mrecalls[1:] != mrecalls[:-1])[0]
    # and sum (\Delta recall) * prec
    ap = np.sum((mrecalls[i + 1] - mrecalls[i]) * mprecisions[i + 1])
    pr_x = np.concatenate(([mrecalls[0]], mrecalls[i + 1]))
    pr_y = np.concatenate(([mprecisions[0]], mprecisions[i + 1]))


    # 绘制ROC曲线，求AUC
    fpr = fp / np.maximum(float(nd), np.finfo(np.float32).eps)
    tpr = tp / np.maximum(float(nd), np.finfo(np.float32).eps)
    mfpr = np.concatenate(([0.], fpr, [1.]))
    mtpr = np.concatenate(([0.], tpr, [1.]))

    i = np.where(mfpr[1:] != mfpr[:-1])[0]
    auc = np.sum((mfpr[i + 1] - mfpr[i]) * mtpr[i + 1])
    roc_x = np.concatenate(([mfpr[0]], mfpr[i + 1]))
    roc_y = np.concatenate(([mtpr[0]], mtpr[i + 1]))



    print('class:{} prec:{:.3f} recall:{:.3f} F1:{:.3f} mAP:{:.3f} AUC:{:.3f} temp_fp:{}\n'.format(
        c, prec_val, rec_val, F1, ap, auc, temp_fp))
    resultoutpath = os.path.join(OUTPUT_PATH, 'predict.txt')
    resulttxt = open(resultoutpath, "a")
    resulttxt.write('class:{} prec:{:.3f} recall:{:.3f} F1:{:.3f} mAP:{:.3f} AUC:{:.3f} temp_fp:{}\n'.format(
        c, prec_val, rec_val, F1, ap, auc, temp_fp))
    resulttxt.close()
    return


# 输出路径
outputpath = os.path.join(OUTPUT_PATH, 'TestCurves')
if not os.path.isdir(outputpath):
    os.makedirs(outputpath)

#类别包括背景 7+1
NUM_CLASSES = 8
TEST_IOU_FILT_TH = 0.5
BACKGROUND_LABEL_ID = 0
TEST_KEEP_TOPK = 25
INPUT_ROIS_PER_IMAGE = 12
file_path = 'E:\\20210823_AnnotatedBreastDatas\\yolo_dataset\\test\\test\\'
f_names = glob.glob(file_path + '*.jpg')
num_test = len(f_names)
conf_thresh = 0.4
nms_thresh = 0.6

# all_boxes:保存所有的预测值
# all_gt_infos: 保存所有的gt值
all_boxes = [[[] for _ in range(num_test)] for _ in range(NUM_CLASSES)]
all_gt_infos = [[[] for _ in range(num_test)] for _ in range(NUM_CLASSES)]


for i in range(len(f_names)):
    # if f_names[i] != '01BB8F2539C341C79DA1D00559786324__LVYvbFEnhzAcPHR2.jpg':
    #     continue
    orig_img = cv2.imdecode(np.fromfile(f_names[i], dtype=np.uint8), 1)
    aaa = orig_img.copy()
    width = aaa.shape[1]
    height = aaa.shape[0]

    # # 读入gt标注
    annotations = np.zeros((INPUT_ROIS_PER_IMAGE, 5), dtype=np.float32)

    gt_file = os.path.join(file_path, "{}.txt".format(f_names[i].split("\\")[-1].split(".jpg")[0]))
    with open(gt_file, "r") as gtf:
        gt_anno = gtf.readlines()
        gt_ind = 0
        for ind in range(len(gt_anno)):
            gt_info = gt_anno[ind].strip().split('\n')
            if len(gt_info[0]) > 0:
                bbox_floats = np.fromstring(gt_info[0], dtype=np.float32, sep=' ')
                label_gt= int(bbox_floats[0]) + 1
                x_yolo_gt = bbox_floats[1]
                y_yolo_gt = bbox_floats[2]
                width_yolo_gt = bbox_floats[3]
                height_yolo_gt = bbox_floats[4]

                top_gt = int((y_yolo_gt - height_yolo_gt / 2) * height)
                left_gt = int((x_yolo_gt - width_yolo_gt / 2) * width)
                right_gt = int((x_yolo_gt + width_yolo_gt / 2) * width)
                bottom_gt = int((y_yolo_gt + height_yolo_gt / 2) * height)
                annotations[gt_ind, :] = [left_gt, top_gt, right_gt, bottom_gt, label_gt]
                gt_ind += 1

    gtf.close()


    for c, _ in enumerate(class_list):
        if c == BACKGROUND_LABEL_ID:
            continue
        cls_gt_boxes = annotations[np.where(annotations[:, -1] == c)]
        gt_box_count = len(cls_gt_boxes)
        if gt_box_count > 0:
            if class_id_map is None:
                all_gt_infos[c][i] = np.reshape(cls_gt_boxes[:, 0:4], gt_box_count * 4).tolist()
            else:
                all_gt_infos[class_id_map[c]][i] = np.reshape( cls_gt_boxes[:, 0:4], gt_box_count * 4).tolist()



    resized = cv2.resize(orig_img, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)
    # img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)
    # img_in = np.expand_dims(img_in, axis=2)
    img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
    img_in = np.expand_dims(img_in, axis=0)
    img_in /= 255.0

    outputs = session.run(None, {input_name: img_in})


    # [batch, num, 1, 4]
    box_array = outputs[0]
    # [batch, num, num_classes]
    confs = outputs[1]

    if type(box_array).__name__ != 'ndarray':
        box_array = box_array.cpu().detach().numpy()
        confs = confs.cpu().detach().numpy()
    num_classes = confs.shape[2]
    # [batch, num, 4]
    box_array = box_array[:, :, 0]

    # [batch, num, num_classes] --> [batch, num]
    max_conf = np.max(confs, axis=2)
    max_id = np.argmax(confs, axis=2)

    assert box_array.shape[0] == 1,"每次单幅图像预测，batch为1"

    argwhere = max_conf[0] > conf_thresh
    l_box_array = box_array[0, argwhere, :]
    l_max_conf = max_conf[0, argwhere]
    l_max_id = max_id[0, argwhere]

    bboxes = []
    # nms for each class
    for j in range(num_classes):
        cls_argwhere = l_max_id == j
        ll_box_array = l_box_array[cls_argwhere, :]
        ll_max_conf = l_max_conf[cls_argwhere]
        ll_max_id = l_max_id[cls_argwhere]

        keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)

        if (keep.size > 0):
            ll_box_array = ll_box_array[keep, :]
            ll_max_conf = ll_max_conf[keep]
            ll_max_id = ll_max_id[keep]

            for k in range(ll_box_array.shape[0]):
                bboxes.append(
                    [
                     ll_max_id[k] + 1,
                     ll_max_conf[k],
                     int(ll_box_array[k, 0] * width),
                     int(ll_box_array[k, 1] * height),
                     int(ll_box_array[k, 2] * width),
                     int(ll_box_array[k, 3] * height),
                     ])

    pred_annos = np.zeros((TEST_KEEP_TOPK, 6))
    valid_box_ind = 0

    if len(bboxes) == 0:
        print(1)
    for t in range(len(bboxes)):
        pred_annos[valid_box_ind, :] = bboxes[t]
        valid_box_ind += 1

    for ind in range(pred_annos.shape[0]):
        c = int(pred_annos[ind][0])
        if c <= 0:
            continue
        score = pred_annos[ind][1]
        box_left = int(pred_annos[ind][2])
        box_top = int(pred_annos[ind][3])
        box_right = int(pred_annos[ind][4])
        box_bottom = int(pred_annos[ind][5])
        if box_right > box_left and box_bottom > box_top:
            if class_id_map is None:
                all_boxes[c][i].append(box_left)
                all_boxes[c][i].append(box_top)
                all_boxes[c][i].append(box_right)
                all_boxes[c][i].append(box_bottom)
                all_boxes[c][i].append(score)
            else:
                all_boxes[class_id_map[c]][i].append(box_left)
                all_boxes[class_id_map[c]][i].append(box_top)
                all_boxes[class_id_map[c]][i].append(box_right)
                all_boxes[class_id_map[c]][i].append(box_bottom)
                all_boxes[class_id_map[c]][i].append(score)

# # 当所有图像都测试完毕
evaluated_class_groups = []
for c, _ in enumerate(class_list):
    if c != BACKGROUND_LABEL_ID:
        if class_id_map is None:
            evaluate_detections(c, all_boxes[c], all_gt_infos[c], outputpath, ovthresh=TEST_IOU_FILT_TH)
        else:
            g = class_id_map[c]
            if g not in evaluated_class_groups:
                evaluate_detections(g, all_boxes[g], all_gt_infos[g], outputpath, ovthresh=TEST_IOU_FILT_TH)
                evaluated_class_groups.append(g)