#!/usr/bin/env python
"""
简化版可视化脚本 - 不依赖复杂库
"""
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import torch

def main():
    # 直接指定结果文件路径
    result_file = '/data/infer_test/20251120_124755/one_batch_results.pkl'

    print(f"加载推理结果: {result_file}")
    with open(result_file, 'rb') as f:
        results = pickle.load(f)

    print(f"结果包含 {len(results)} 个样本")

    # 查看第一个样本
    sample = results[0]
    print(f"样本 0 的keys: {list(sample.keys())}")

    # 1. 分析BEV分割
    if 'masks_bev' in sample:
        masks_bev = sample['masks_bev']
        print(f"BEV分割形状: {masks_bev.shape}")

        # 转换为numpy
        if torch.is_tensor(masks_bev):
            masks_bev = masks_bev.cpu().numpy()

        # 为每个类别创建可视化
        class_names = ['drivable_area', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area', 'divider']

        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        fig.suptitle('BEV分割结果 - 第一个样本', fontsize=16)

        for i in range(6):
            row, col = i // 3, i % 3
            ax = axes[row, col]

            mask = masks_bev[i]
            ax.imshow(mask, cmap='gray', vmin=0, vmax=1)
            ax.set_title(f'{class_names[i]}\n像素占比: {(mask > 0.5).sum() / mask.size * 100:.1f}%')
            ax.axis('off')

        plt.tight_layout()
        plt.savefig('/data/infer_test/bev_segmentation_sample_0.png', dpi=150, bbox_inches='tight')
        print("BEV分割可视化已保存: /data/infer_test/bev_segmentation_sample_0.png")
        plt.show()

        # 打印分割统计
        print("\nBEV分割统计:")
        for i, name in enumerate(class_names):
            pixel_count = (masks_bev[i] > 0.5).sum()
            percentage = pixel_count / masks_bev[i].size * 100
            print("20s")

    # 2. 分析3D检测
    if 'boxes_3d' in sample:
        boxes_3d = sample['boxes_3d']
        scores_3d = sample['scores_3d']
        labels_3d = sample['labels_3d']

        print(f"\n3D检测结果:")
        print(f"检测框数量: {len(boxes_3d)}")

        if len(scores_3d) > 0:
            if torch.is_tensor(scores_3d):
                scores_3d = scores_3d.cpu().numpy()
            print(".3f")
            print(f"检测类别分布: {np.bincount(labels_3d, minlength=10) if len(labels_3d) > 0 else '无检测'}")

    # 3. 分析Ground Truth (如果有)
    if 'gt_masks_bev' in sample:
        gt_masks = sample['gt_masks_bev']
        print(f"\nGround Truth BEV分割形状: {gt_masks.shape}")

        # 计算预测与GT的差异
        if 'masks_bev' in sample:
            pred_masks = sample['masks_bev']
            if torch.is_tensor(pred_masks):
                pred_masks = pred_masks.cpu().numpy()
            if torch.is_tensor(gt_masks):
                gt_masks = gt_masks.cpu().numpy()

            print("\n预测 vs Ground Truth 对比:")
            for i, name in enumerate(class_names):
                pred_pixels = (pred_masks[i] > 0.5).sum()
                gt_pixels = (gt_masks[i] > 0.5).sum()
                iou = np.logical_and(pred_masks[i] > 0.5, gt_masks[i] > 0.5).sum() / max(
                    np.logical_or(pred_masks[i] > 0.5, gt_masks[i] > 0.5).sum(), 1)
                print("20s")

if __name__ == '__main__':
    main()