import matplotlib.pyplot as plt
from detectron2.data import build_detection_train_loader, DatasetCatalog, MetadataCatalog
from detectron2.utils.visualizer import Visualizer
from detrex.config import get_config

# 假设 dataloader 已经从配置文件中导入
dataloader = get_config("common/data/coco_ovd.py").dataloader

# 可视化样本
def visualize_samples(dataloader_cfg, num_samples=5):
    # 构建训练数据加载器
    train_loader = build_detection_train_loader(dataloader_cfg.train)

    for idx, batch in enumerate(train_loader):
        if idx >= num_samples:
            break

        sample = batch[0]
        image = sample["image"].permute(1, 2, 0).cpu().numpy()
        annotations = sample.get("instances", None)

        print(f"Sample {idx + 1}:")
        print(f"  Image shape: {image.shape}")
        print(f"  Annotations: {annotations}")

        # 可视化
        visualizer = Visualizer(image, metadata=MetadataCatalog.get("coco_2017_train_ovd_unipro"))
        vis = visualizer.draw_instance_predictions(annotations)

        plt.figure(figsize=(12, 8))
        plt.imshow(vis.get_image())
        plt.axis("off")
        plt.title(f"Sample {idx + 1}")
        plt.show()

# 调用函数可视化前 5 个样本
visualize_samples(dataloader, num_samples=5)
