diffusion-dpo-coco-ade/verify_annotations.py

"""
验证脚本：检查 annotations 是否正确，以及在原始 COCO 上测试检测器
"""

import os
import json
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image
import numpy as np

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN_ResNet50_FPN_V2_Weights
from torchvision.transforms import functional as TF
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval


# ============================================================================
# 配置
# ============================================================================
CONFIG = {
    # 原始 COCO 数据
    'original_coco_dir': '/home/wanghongbo06/baipurui/DATA/COCO/val2017',
    'original_ann_file': '/home/wanghongbo06/baipurui/DATA/COCO/annotations/instances_val2017.json',
    
    # prepare 生成的数据
    'patch_gt_dir': '/home/wanghongbo06/baipurui/DATA/COCO_patch/gt',
    'patch_ann_file': '/home/wanghongbo06/baipurui/DATA/COCO_patch/patch_annotations.json',
    
    'device': 'cuda',
}
# ============================================================================

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


def visualize_annotations(ann_file, image_dir, num_images=5):
    """可视化 GT annotations"""
    print("="*60)
    print("检查 Annotations 格式")
    print("="*60)
    
    coco = COCO(ann_file)
    image_ids = list(coco.imgs.keys())[:num_images]
    
    fig, axes = plt.subplots(1, num_images, figsize=(4*num_images, 4))
    if num_images == 1:
        axes = [axes]
    
    for ax, img_id in zip(axes, image_ids):
        img_info = coco.imgs[img_id]
        img_path = Path(image_dir) / img_info['file_name']
        
        if not img_path.exists():
            print(f"图片不存在: {img_path}")
            continue
        
        img = Image.open(img_path)
        ax.imshow(img)
        ax.set_title(f"ID: {img_id}\n{img_info['file_name']}")
        
        # 绘制 GT bbox
        ann_ids = coco.getAnnIds(imgIds=img_id)
        anns = coco.loadAnns(ann_ids)
        
        for ann in anns:
            bbox = ann['bbox']  # [x, y, w, h]
            rect = patches.Rectangle(
                (bbox[0], bbox[1]), bbox[2], bbox[3],
                linewidth=2, edgecolor='r', facecolor='none'
            )
            ax.add_patch(rect)
        
        ax.axis('off')
        print(f"Image {img_id}: {len(anns)} annotations")
    
    plt.tight_layout()
    plt.savefig('verify_annotations.png', dpi=150)
    print(f"\n可视化保存到: verify_annotations.png")


def check_annotation_format(ann_file):
    """检查 annotation 格式"""
    print("\n" + "="*60)
    print("检查 Annotation JSON 格式")
    print("="*60)
    
    with open(ann_file, 'r') as f:
        data = json.load(f)
    
    print(f"Images 数量: {len(data['images'])}")
    print(f"Annotations 数量: {len(data['annotations'])}")
    print(f"Categories 数量: {len(data['categories'])}")
    
    # 检查 image ids
    image_ids = [img['id'] for img in data['images']]
    print(f"\nImage ID 范围: {min(image_ids)} - {max(image_ids)}")
    print(f"Image ID 是否连续: {image_ids == list(range(len(image_ids)))}")
    
    # 检查 annotation 的 image_id 是否都有对应的 image
    ann_image_ids = set(ann['image_id'] for ann in data['annotations'])
    missing = ann_image_ids - set(image_ids)
    if missing:
        print(f"警告: {len(missing)} 个 annotation 的 image_id 找不到对应的 image!")
    else:
        print("所有 annotation 的 image_id 都有对应的 image ✓")
    
    # 检查 category_ids
    ann_cat_ids = set(ann['category_id'] for ann in data['annotations'])
    print(f"\n使用的 category_id: {sorted(ann_cat_ids)[:10]}... (共 {len(ann_cat_ids)} 个)")
    
    # 检查第一个 annotation
    if data['annotations']:
        print("\n第一个 annotation 示例:")
        ann = data['annotations'][0]
        for key, value in ann.items():
            if key == 'segmentation':
                print(f"  {key}: {type(value)}, 长度={len(value) if value else 0}")
            else:
                print(f"  {key}: {value}")


def test_on_original_coco(coco_dir, ann_file, device='cuda', num_images=100):
    """在原始 COCO 上测试 Mask R-CNN"""
    print("\n" + "="*60)
    print("在原始 COCO 上测试 Mask R-CNN")
    print("="*60)
    
    coco = COCO(ann_file)
    image_ids = list(coco.imgs.keys())[:num_images]
    
    print(f"测试图片数: {num_images}")
    
    # 加载模型
    print("加载 Mask R-CNN...")
    weights = MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT
    model = maskrcnn_resnet50_fpn_v2(weights=weights)
    model.eval()
    model.to(device)
    
    # 注意：torchvision 的 Mask R-CNN 输出的 labels 已经是 COCO category_id
    # 不需要额外映射！
    
    bbox_results = []
    
    print("运行检测...")
    from tqdm import tqdm
    for img_id in tqdm(image_ids):
        img_info = coco.imgs[img_id]
        img_path = Path(coco_dir) / img_info['file_name']
        
        if not img_path.exists():
            continue
        
        img = Image.open(img_path).convert('RGB')
        img_tensor = TF.to_tensor(img).unsqueeze(0).to(device)
        
        with torch.no_grad():
            outputs = model(img_tensor)[0]
        
        boxes = outputs['boxes'].cpu().numpy()
        labels = outputs['labels'].cpu().numpy()
        scores = outputs['scores'].cpu().numpy()
        
        for i in range(len(boxes)):
            if scores[i] < 0.05:  # 提高阈值过滤低置信度
                continue
            
            x1, y1, x2, y2 = boxes[i]
            category_id = int(labels[i])  # 直接使用，已经是 COCO category_id
            
            bbox_results.append({
                'image_id': img_id,
                'category_id': category_id,
                'bbox': [float(x1), float(y1), float(x2-x1), float(y2-y1)],
                'score': float(scores[i]),
            })
    
    print(f"检测到 {len(bbox_results)} 个物体")
    
    # 评估
    if bbox_results:
        coco_dt = coco.loadRes(bbox_results)
        coco_eval = COCOeval(coco, coco_dt, 'bbox')
        coco_eval.params.imgIds = image_ids
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()
        
        print(f"\n原始 COCO 上的 AP: {coco_eval.stats[0]*100:.2f}%")
        return coco_eval.stats[0]
    else:
        print("没有检测到任何物体!")
        return 0


def main():
    print("="*60)
    print("Annotations 验证工具")
    print("="*60)
    
    # 1. 检查 annotation 格式
    if Path(CONFIG['patch_ann_file']).exists():
        check_annotation_format(CONFIG['patch_ann_file'])
    
    # 2. 可视化 annotations
    if Path(CONFIG['patch_gt_dir']).exists() and Path(CONFIG['patch_ann_file']).exists():
        visualize_annotations(
            CONFIG['patch_ann_file'],
            CONFIG['patch_gt_dir'],
            num_images=5
        )
    
    # 3. 在原始 COCO 上测试（作为 baseline）
    if Path(CONFIG['original_coco_dir']).exists():
        device = CONFIG['device']
        if device == 'cuda' and not torch.cuda.is_available():
            device = 'cpu'
        
        ap = test_on_original_coco(
            CONFIG['original_coco_dir'],
            CONFIG['original_ann_file'],
            device=device,
            num_images=100  # 测试 100 张原图
        )
        
        print("\n" + "="*60)
        print("结论")
        print("="*60)
        if ap > 0.3:
            print(f"原始 COCO 上 AP={ap*100:.1f}%，检测器正常")
            print("问题可能在于 resize 后的 annotations 转换")
        else:
            print(f"原始 COCO 上 AP={ap*100:.1f}%，检测器或评估代码可能有问题")


if __name__ == '__main__':
    main()