对目标检测模型 PR 曲线的原理分析和实现

pred_boxes: Boxes = ...
gt_boxes: Boxes = ...
iou_matrix: torch.Tensor = pairewise_iou(pred_boxes, gt_boxes)
ious, gt_ids = iou_matrix.max(dim=1)

df_i = pd.DataFrame({
        "gt_id": [f"{image_id}_{i}" for i in gt_ids.cpu().numpy().tolist()],
        "predict_id": [f"{image_id}_{i}" for i in range(len(gt_ids))],
        "iou": ious.cpu().numpy().tolist(),
        "confidence": output["instances"].scores.cpu().numpy().tolist()
    }

df=pd.concat((df, df_i), axis=0)

df["TP"] = df["iou"] > 0.5

max_iou=df[["gt_id", "iou"]].groupby("gt_id").max()
best_match_df=df[["gt_id", "iou", "predict_id"]].merge(max_iou, how="inner", on = ["gt_id", "iou"])
best_match_df=best_match_df.drop(columns=["iou"])
best_match_df["best_match"]=pd.Series([True for _ in range(len(best_match_df))])

df = df.merge(best_match_df, how="outer", on=["gt_id", "predict_id"]).fillna(False)

df["TP"] = df["TP"] & df["best_match"]

df = df.sort_values(by="confidence", ascending=False)

M = pd.Series([i + 1 for i in range(len(df))])

P = pd.Series([tp[0:i+1].sum() for i in range(len(df))])

precise = P / M
recall = P / N

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

from concern.config import Configurable,Config
import numpy as np
import math

import torch
from torchvision.io import read_image
from torchvision.io import ImageReadMode
import torch.nn.functional as F
import glob
import math
import cv2
from shapely.geometry import Polygon
import pandas as pd

def load_eval_model():
    """
    加载模型
    """

    args={
        "resume": "./checkpoints/td500_resnet50",
        "box_thresh": 0.4, 
        # "thresh": 0.8,
        "log_dir": "./logs", 
        "visualize": True,
        "image_short_side": 736, 
        "polygon": False, 
        "result_dir": "./result_dir"}

    config = Config()
    experiment_args=config.compile(config.load("./experiments/seg_detector/td500_resnet50_deform_thre.yaml"))['Experiment']
    experiment_args.update(cmd=args)
    experiment = Configurable.construct_class_from_config(experiment_args)

    experiment.load("evaluation", **experiment_args)

    structure = experiment.structure
    model_path = args["resume"]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = structure.builder.build(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

    return model, structure


def load_image(path, short_side=736):
    """
    加载图片
    """
    RGB_MEAN = torch.Tensor([122.67891434, 116.66876762, 104.00698793]).view(-1, 1, 1)
    img = read_image(path, ImageReadMode.RGB)
    
    c, h, w = img.shape
    original_size = (h, w)
    if h < w:
        new_height = short_side
        new_width = int(math.ceil(new_height / h * w / 32) * 32)
    else:
        new_width = short_side
        new_height = int(math.ceil(new_width / w * h / 32) * 32)

    img = img.unsqueeze(0).float()
    img = F.interpolate(img, size=(new_height, new_width), mode='bilinear', align_corners=False)
    img = img - RGB_MEAN
    img /= 255.0

    return img, original_size

def read_ground_truth(path):
    """
    读取 ground truth 数据

    :param path: ground truth 文件路径
    """
    with open(path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip() for line in lines]
    gt_boxes = []
    for line in lines:
        index, diffcult, x, y, w, h, theta = line.split(' ')
        x, y, w, h, theta = float(x), float(y), float(w), float(h), float(theta)
        cx = x + w / 2
        cy = y + h / 2
        x1 = cx - w / 2 * math.cos(theta) - h / 2 * math.sin(theta)
        y1 = cy - w / 2 * math.sin(theta) + h / 2 * math.cos(theta)
        x2 = cx + w / 2 * math.cos(theta) - h / 2 * math.sin(theta)
        y2 = cy + w / 2 * math.sin(theta) + h / 2 * math.cos(theta)
        x3 = cx + w / 2 * math.cos(theta) + h / 2 * math.sin(theta)
        y3 = cy + w / 2 * math.sin(theta) - h / 2 * math.cos(theta)
        x4 = cx - w / 2 * math.cos(theta) + h / 2 * math.sin(theta)
        y4 = cy - w / 2 * math.sin(theta) - h / 2 * math.cos(theta)
        x1, y1, x2, y2, x3, y3, x4, y4 = int(x1), int(y1), int(x2), int(y2), int(x3), int(y3), int(x4), int(y4)
        gt_boxes.append([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])
    return gt_boxes # n x 4 x 2

def pairwise_iou(gt_boxes, pred_boxes):
    """
    计算 gt 列表与 pred 列表之间的 IOU 矩阵

    :params gt_boxes: gt 列表 [n, 4, 2]
    :params pred_boxes: pred 列表 [m, 4, 2]
    """

    iou_matrix = np.zeros((len(pred_boxes), len(gt_boxes)))
    for i in range(len(pred_boxes)):
        for j in range(len(gt_boxes)):
            gt_polygon = Polygon(gt_boxes[j])
            polygon =  Polygon(pred_boxes[i])
            iou_matrix[i, j] = polygon.intersection(gt_polygon).area / polygon.union(gt_polygon).area

    return iou_matrix # m x n

test_dir = "./MSRA-TD500/test"
img_paths = glob.glob(test_dir + "/*.jpg")

df = pd.DataFrame()
N = 0 # 所有图片上的文本实例数量

for img_path in img_paths:

    filename=os.path.basename(img_path)
    gt_path=test_dir + "/" + os.path.splitext(filename)[0] + ".gt"
    gt_boxes = read_ground_truth(gt_path)

    img, original_size=load_image(img_path)
    batch = {
        "filename": [img_path],
        "shape": [original_size],
        "image": img
    }
    
    with torch.no_grad():
        pred = model.forward(batch, training = False)
        output = structure.representer.represent(batch, pred, is_output_polygon=False)
        a, b = output
        
        pred_boxes = a[0]
        confidence = b[0]
        pred_boxes=pred_boxes.tolist()
    
    N += len(gt_boxes)

    if len(pred_boxes) == 0 or len(gt_boxes) == 0:
        continue

    iou_matrix=pairwise_iou(gt_boxes, pred_boxes)
    gt_ids = iou_matrix.argmax(axis=1)
    ious=iou_matrix[range(len(iou_matrix)), gt_ids]

    df_i = pd.DataFrame(
        {
            "gt_id": [f"{filename}_{i}" for i in gt_ids], 
            "pred_id": [f"{filename}_{i}" for i in range(len(pred_boxes))], 
            "iou": ious,
            "confidence": confidence
        }
    )

    df = pd.concat([df, df_i], axis=0)
    

df["TP"] = df["iou"] > 0.5  # 以 0.5 为 IOU 阈值设置 True Positive

## 解决多个 gt 对应同一个 pred 的问题
max_iou=df[["gt_id", "iou"]].groupby("gt_id").max()
best_match_df=df[["gt_id", "iou", "pred_id"]].merge(max_iou, how="inner", on = ["gt_id", "iou"])
best_match_df=best_match_df.drop(columns=["iou"])
best_match_df["best_match"]=pd.Series([True for _ in range(len(best_match_df))])

df = df.merge(best_match_df, how="outer", on=["gt_id", "pred_id"]).fillna(False)

## 解决 IOU 大于阈值但不是最佳匹配的对
df["TP"] = df["TP"] & df["best_match"]

## 对 confidence 按从大到小排序
df = df.sort_values(by="confidence", ascending=False) 

## 大于各个 confidence 的预测数量
M = pd.Series([i + 1 for i in range(len(df))])

## 大于各个 confidence 的正预测数量
tp = df["TP"]
P = pd.Series([tp[0:i+1].sum() for i in range(len(df))])

precision = P / M
recall = P / N