import torch
import os
from pathlib import Path

import os
os.environ['GITHUB_ASSETS'] = 'disabled'
os.environ['YOLOv5_OFFLINE'] = 'true'
os.environ['ULTRALYTICS_HUB'] = 'off'

# 强制禁用所有网络连接
os.environ['GITHUB_ASSETS'] = 'disabled'
os.environ['YOLOv5_OFFLINE'] = 'true'

# 配置路径
model_dir = Path('/home/th/jcq/AI_AutoPlat/AI_web_dsj/ultralytics/yolov5')
model_path = Path('/home/th/jcq/AI_AutoPlat/yolov5-th/yolov5/yolov5s.pt')  # 自定义模型权重路径
device = torch.device('cpu')  # 离线环境建议强制使用CPU

def load_yolov5_offline():
    try:
        # 1. 手动添加本地YOLOv5到Python路径
        import sys
        sys.path.insert(0, str(model_dir))
        
        # 2. 直接使用YOLOv5本地代码加载
        from models.experimental import attempt_load
        model = attempt_load(
            weights=model_path,
            device=device,
            inplace=True,
            fuse=True  # 融合Conv+BN层提升效率
        )
        
        # 3. 验证模型
        if not hasattr(model, 'names'):
            raise RuntimeError("模型结构异常")
            
        print(f"✅ 离线加载成功 | 设备: {device} | 类别数: {len(model.names)}")
        return model
        
    except Exception as e:
        print(f"❌ 加载失败: {type(e).__name__}: {e}")
        if "No module named" in str(e):
            print("\n⚠️ 解决方案:")
            print(f"1. 确认 {model_dir} 包含完整YOLOv5代码")
            print(f"2. 检查是否有 __init__.py 文件在 models/ 和 utils/ 目录")
        return None

if __name__ == "__main__":
    model = load_yolov5_offline()
    if model:
        # 测试推理
        img = torch.rand((1, 3, 640, 640))  # 模拟输入图像
        results = model(img)
        print(f"推理测试完成！输出形状: {results[0].shape}")