55 lines
1.9 KiB
Python
Executable File
55 lines
1.9 KiB
Python
Executable File
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}") |