test_ap_on_coco_onnx.py

import argparse
import os
import sys
import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, DistributedSampler
import torchvision
import onnxruntime as ort

import groundingdino.datasets.transforms as T
from groundingdino.util import box_ops, get_tokenlizer
from groundingdino.util.misc import collate_fn
from groundingdino.util.slconfig import SLConfig
from groundingdino.util.vl_utils import build_captions_and_token_span, create_positive_map_from_span
from groundingdino.datasets.cocogrounding_eval import CocoGroundingEvaluator


class CocoDetection(torchvision.datasets.CocoDetection):
    def __init__(self, img_folder, ann_file, transforms):
        super().__init__(img_folder, ann_file)
        self._transforms = transforms

    def __getitem__(self, idx):
        img, target = super().__getitem__(idx)

        w, h = img.size
        boxes = [obj["bbox"] for obj in target]
        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
        boxes[:, 2:] += boxes[:, :2]  # xywh -> xyxy
        boxes[:, 0::2].clamp_(min=0, max=w)
        boxes[:, 1::2].clamp_(min=0, max=h)
        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
        boxes = boxes[keep]

        target_new = {}
        image_id = self.ids[idx]
        target_new["image_id"] = image_id
        target_new["boxes"] = boxes
        target_new["orig_size"] = torch.as_tensor([int(h), int(w)])

        if self._transforms is not None:
            img, target = self._transforms(img, target_new)

        return img, target


class PostProcessCocoGrounding(nn.Module):
    """保持和原代码一致的后处理逻辑"""
    def __init__(self, num_select=300, coco_api=None, tokenlizer=None) -> None:
        super().__init__()
        self.num_select = num_select

        assert coco_api is not None
        category_dict = coco_api.dataset['categories']
        cat_list = [item['name'] for item in category_dict]
        captions, cat2tokenspan = build_captions_and_token_span(cat_list, True)
        tokenspanlist = [cat2tokenspan[cat] for cat in cat_list]
        positive_map = create_positive_map_from_span(
            tokenlizer(captions), tokenspanlist)  # 80, 256. normed

        id_map = {0: 1, 1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: 7, 7: 8, 8: 9, 9: 10, 10: 11, 11: 13, 12: 14, 13: 15, 14: 16, 15: 17, 16: 18, 17: 19, 18: 20, 19: 21, 20: 22, 21: 23, 22: 24, 23: 25, 24: 27, 25: 28, 26: 31, 27: 32, 28: 33, 29: 34, 30: 35, 31: 36, 32: 37, 33: 38, 34: 39, 35: 40, 36: 41, 37: 42, 38: 43, 39: 44, 40: 46,
                  41: 47, 42: 48, 43: 49, 44: 50, 45: 51, 46: 52, 47: 53, 48: 54, 49: 55, 50: 56, 51: 57, 52: 58, 53: 59, 54: 60, 55: 61, 56: 62, 57: 63, 58: 64, 59: 65, 60: 67, 61: 70, 62: 72, 63: 73, 64: 74, 65: 75, 66: 76, 67: 77, 68: 78, 69: 79, 70: 80, 71: 81, 72: 82, 73: 84, 74: 85, 75: 86, 76: 87, 77: 88, 78: 89, 79: 90}

        new_pos_map = torch.zeros((91, 256))
        for k, v in id_map.items():
            new_pos_map[v] = positive_map[k]
        self.positive_map = new_pos_map

    @torch.no_grad()
    def forward(self, outputs, target_sizes, not_to_xyxy=False):
        num_select = self.num_select
        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']

        prob_to_token = torch.from_numpy(out_logits).sigmoid()  # 适配numpy输入
        pos_maps = self.positive_map.to(prob_to_token.device)
        prob_to_label = prob_to_token @ pos_maps.T

        assert prob_to_label.shape[0] == len(target_sizes)
        assert target_sizes.shape[1] == 2

        prob = prob_to_label
        topk_values, topk_indexes = torch.topk(
            prob.view(prob_to_label.shape[0], -1), num_select, dim=1)
        scores = topk_values
        topk_boxes = topk_indexes // prob.shape[2]
        labels = topk_indexes % prob.shape[2]

        if not_to_xyxy:
            boxes = torch.from_numpy(out_bbox)
        else:
            boxes = box_ops.box_cxcywh_to_xyxy(torch.from_numpy(out_bbox))

        boxes = torch.gather(
            boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))

        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes = boxes * scale_fct[:, None, :]

        results = [{'scores': s, 'labels': l, 'boxes': b}
                   for s, l, b in zip(scores, labels, boxes)]

        return results


def load_onnx_model(onnx_path, device="cuda"):
    """加载ONNX模型并创建推理session"""
    providers = ['CPUExecutionProvider']
    if device == "cuda" and ort.get_device() == "GPU":
        providers = ['ROCMExecutionProvider', 'CPUExecutionProvider']
    
    session = ort.InferenceSession(
        onnx_path,
        providers=providers,
        provider_options=[{'device_id': 0}] if "CUDAExecutionProvider" in providers else []
    )
    return session


def onnx_inference(session, images, captions):
    """ONNX模型推理（需匹配模型输入格式）"""
    # 转换为numpy（ONNX Runtime不支持torch tensor）
    images_np = images.cpu().numpy().astype(np.float32)
    
    # 注意：此处需根据你的ONNX模型输入名调整（可通过netron查看）
    # 假设模型输入为 "images" 和 "captions"（需根据实际情况修改）
    input_feed = {
        session.get_inputs()[0].name: images_np,
        # 如果caption是文本token，需补充token化逻辑，此处假设已处理
        # session.get_inputs()[1].name: captions_np
    }
    
    # 执行推理
    outputs = session.run(None, input_feed)
    
    # 解析输出（需匹配模型输出格式，假设输出为logits和bbox）
    # 需根据你的ONNX模型输出调整维度和顺序
    pred_logits = outputs[0]  # 形状: [bs, 100, 256]
    pred_boxes = outputs[1]   # 形状: [bs, 100, 4]
    
    return {"pred_logits": pred_logits, "pred_boxes": pred_boxes}


def main(args):
    # 配置（主要用于tokenizer和后处理）
    cfg = SLConfig.fromfile(args.config_file)

    # 加载ONNX模型
    onnx_session = load_onnx_model(args.onnx_path, args.device)

    # 构建数据加载器（和原代码一致）
    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    dataset = CocoDetection(
        args.image_dir, args.anno_path, transforms=transform)
    data_loader = DataLoader(
        dataset, batch_size=1, shuffle=False, num_workers=args.num_workers, collate_fn=collate_fn)

    # 构建后处理器（和原代码一致）
    tokenlizer = get_tokenlizer.get_tokenlizer(cfg.text_encoder_type)
    postprocessor = PostProcessCocoGrounding(
        coco_api=dataset.coco, tokenlizer=tokenlizer)

    # 构建评估器（和原代码一致）
    evaluator = CocoGroundingEvaluator(
        dataset.coco, iou_types=("bbox",), useCats=True)

    # 构建文本提示（和原代码一致）
    category_dict = dataset.coco.dataset['categories']
    cat_list = [item['name'] for item in category_dict]
    caption = " . ".join(cat_list) + ' .'
    print("Input text prompt:", caption)

    # 运行推理
    start = time.time()
    for i, (images, targets) in enumerate(data_loader):
        # 预处理图像（和原代码一致）
        images = images.tensors.to(args.device)
        bs = images.shape[0]
        input_captions = [caption] * bs

        # ONNX推理（替换原PyTorch模型推理）
        outputs = onnx_inference(onnx_session, images, input_captions)

        # 后处理（适配ONNX输出格式）
        orig_target_sizes = torch.stack(
            [t["orig_size"] for t in targets], dim=0).to(args.device)
        results = postprocessor(outputs, orig_target_sizes)
        cocogrounding_res = {
            target["image_id"]: output for target, output in zip(targets, results)}
        evaluator.update(cocogrounding_res)

        # 打印进度
        if (i+1) % 30 == 0:
            used_time = time.time() - start
            eta = len(data_loader) / (i+1e-5) * used_time - used_time
            print(
                f"processed {i}/{len(data_loader)} images. time: {used_time:.2f}s, ETA: {eta:.2f}s")

    # 评估指标汇总
    evaluator.synchronize_between_processes()
    evaluator.accumulate()
    evaluator.summarize()

    print("Final results:", evaluator.coco_eval["bbox"].stats.tolist())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        "Grounding DINO ONNX eval on COCO", add_help=True)
    # 新增ONNX模型路径参数
    parser.add_argument("--onnx_path", type=str, required=True, help="path to onnx model file")
    # 保留原配置文件参数（用于tokenizer和后处理）
    parser.add_argument("--config_file", "-c", type=str,
                        required=True, help="path to config file")
    parser.add_argument("--device", type=str, default="cuda",
                        help="running device (default: cuda)")
    # 后处理参数
    parser.add_argument("--num_select", type=int, default=300,
                        help="number of topk to select")
    # COCO数据集参数
    parser.add_argument("--anno_path", type=str,
                        required=True, help="coco annotation path")
    parser.add_argument("--image_dir", type=str,
                        required=True, help="coco image dir")
    parser.add_argument("--num_workers", type=int, default=4,
                        help="number of workers for dataloader")
    args = parser.parse_args()

    main(args)