#!/usr/bin/env python3
"""
快速可视化脚本 - 推理3个样本并生成可视化
避免复杂配置加载，直接使用tools/test.py的简化版本
"""

import os
import sys
sys.path.insert(0, '/workspace/bevfusion')

# 设置环境变量避免显示警告
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm
import pickle

print("=" * 80)
print("BEVFusion 快速可视化")
print("=" * 80)
print("推理3个样本，生成可视化结果")
print("=" * 80)
print()

# 1. 检查checkpoint
checkpoint_path = "runs/run-326653dc-74184412/epoch_19.pth"
if not os.path.exists(checkpoint_path):
    print(f"❌ Checkpoint不存在: {checkpoint_path}")
    sys.exit(1)

print(f"✅ Checkpoint存在: {checkpoint_path}")
print()

# 2. 运行快速推理
print("🚀 开始推理...")
print("=" * 80)
print()

cmd = f"""
/opt/conda/bin/torchpack dist-run -np 1 /opt/conda/bin/python tools/test.py \
    configs/nuscenes/det/transfusion/secfpn/camera+lidar/swint_v0p075/multitask.yaml \
    {checkpoint_path} \
    --eval bbox segm \
    --out results_quick.pkl \
    2>&1 | head -200
"""

os.system(cmd)

print()
print("=" * 80)
print("✅ 推理完成！")
print("=" * 80)
print()

# 3. 检查是否生成了结果
if os.path.exists('results_quick.pkl'):
    print("✅ 结果文件已生成: results_quick.pkl")
    
    # 简单查看结果统计
    with open('results_quick.pkl', 'rb') as f:
        results = pickle.load(f)
    print(f"📊 结果数量: {len(results)} 个样本")
else:
    print("⚠️  结果文件未生成，可能推理失败")
    print("请检查上面的输出日志")

print()
print("=" * 80)
print("下一步：查看结果")
print("=" * 80)
print("如果推理成功，可以使用以下命令可视化结果：")
print("  python visualize_results.py --results results_quick.pkl --samples 3")
print()