nvidia-smi显示的是驱动支持的最高CUDA版本。
nvcc --version显示的是实际安装的CUDA编译器版本
nvidia-smi
labelimg==1.8.6
conda create -n t251p310 python=3.10
conda create -n yolo11 python=3.9
import torch
# 检查PyTorch是否可以使用CUDA
if torch.cuda.is_available():
print("PyTorch supports GPU!")
# 打印出可用的GPU数量
print("Number of GPUs available:", torch.cuda.device_count())
# 打印当前使用的GPU
print("Current GPU:", torch.cuda.current_device())
# 打印GPU的名称
print("GPU name:", torch.cuda.get_device_name(torch.cuda.current_device()))
else:
print("PyTorch does not support GPU.")
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
https://download.pytorch.org/whl/cu124/torch-2.5.1%2Bcu124-cp310-cp310-win_amd64.whl
pip install opencv-python==4.11.0.86 matplotlib numpy==1.26.3 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install opencv-python==4.11.0.86 matplotlib numpy==1.26.3
pip install attrs==19.1.0
pip install dill attrs==19.1.0
from ultralytics import YOLO
import cv2
# Load a model
model = YOLO('yolo11n.pt')
# Run batched inference on a list of images
results = model.predict("bus.jpg", imgsz=640, save=True, device=0,conf=0.8) #gpu
# results = model.predict("bus.jpg", imgsz=640, save=True, device='cpu')
# 获取检测结果图像
result_image = results[0].plot() # 绘制检测结果图像
# 显示检测结果图像
cv2.imshow('Detection Result', result_image)
# 等待按键按下
cv2.waitKey(0)
# 关闭显示窗口
cv2.destroyAllWindows()
# 保存检测后的图像为别名 bus_dec.jpg
cv2.imwrite('bus_dec.jpg', result_image)
根据你当前使用的 ultralytics YOLOE 模型,model.predict 方法常见参数如下(以 ultralytics 8.x/YOLOE 为例):
source:输入图片路径、文件夹、视频路径、摄像头编号等。
conf:置信度阈值(如 0.25),低于该值的检测结果会被过滤。
iou:NMS 的 IoU 阈值(如 0.7)。
classes:只检测指定类别(类别索引列表)。
save:是否保存结果(True/False)。
show:是否显示结果窗口(True/False)。
device:推理设备(如 'cpu'、'cuda:0')。
half:是否使用半精度(True/False)。
imgsz:输入图片尺寸(如 640)。
max_det:每张图片最多检测目标数。
stream:是否流式处理(True/False)。
示例用法:
results = model.predict(
source="g:/Red.jpg",
conf=0.2,
iou=0.5,
imgsz=640,
device='cuda:0',
show=True
)
import cv2
import time
from ultralytics import YOLO
# 手动设置设备 (cuda 或 cpu)
device = 'cuda' # 修改为 'cpu' 使用 CPU
# device = 'cpu'
# 初始化YOLO模型
model = YOLO('yolo11n.pt')
# 将模型移动到指定设备
model.to(device)
# 创建类别映射字典
class_names = {
0: 'person',
2: 'car',
3: 'motorbike',
5: 'bus',
7: 'truck',
# 这里可以继续添加其他类别的映射
}
# 打开视频文件或摄像头
cap = cv2.VideoCapture('nfs.mp4') # 替换为 0 使用摄像头
if not cap.isOpened():
print("Error: Could not open video.")
exit()
# 获取视频帧率和尺寸
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# 初始化视频写入器
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('output.avi', fourcc, fps, (width, height))
while True:
ret, frame = cap.read()
if not ret:
break
start_time = time.time()
# YOLO进行目标检测
results = model(frame) # 使用模型进行推理
detections = results[0].boxes.xyxy.cpu().numpy() # 获取检测结果并转换为numpy数组
confidences = results[0].boxes.conf.cpu().numpy() # 获取置信度
class_ids = results[0].boxes.cls.cpu().numpy() # 获取类别ID
# 绘制检测框和类别
for det, conf, cls in zip(detections, confidences, class_ids):
x1, y1, x2, y2 = map(int, det)
label = class_names.get(int(cls), f'Class {int(cls)}') # 获取类别名称,默认显示为 Class <类别编号>
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f'{label} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# 计算FPS并显示在窗口上
end_time = time.time()
fps_display = 1 / (end_time - start_time)
cv2.putText(frame, f'FPS: {fps_display:.2f}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
# 显示结果
cv2.imshow('Video', frame)
# 写入视频文件
# out.write(frame)
# if cv2.waitKey(1) & 0xFF == ord('q'):
# break
key = cv2.waitKey(1) # 等待用户按键
if key == 27: # 如果用户按下ESC键,退出程序
break
cap.release()
out.release()
cv2.destroyAllWindows()
certifi==2024.8.30
charset-normalizer==3.3.2
colorama==0.4.6
contourpy==1.3.0
cycler==0.12.1
dill==0.3.9
filelock==3.13.1
fonttools==4.54.1
fsspec==2024.2.0
idna==3.10
Jinja2==3.1.3
kiwisolver==1.4.7
MarkupSafe==2.1.5
matplotlib==3.9.2
mpmath==1.3.0
networkx==3.2.1
packaging==24.1
pandas==2.2.3
pillow==10.2.0
psutil==6.0.0
py-cpuinfo==9.0.0
pyparsing==3.1.4
python-dateutil==2.9.0.post0
pytz==2024.2
PyYAML==6.0.2
requests==2.32.3
scipy==1.13.1
seaborn==0.13.2
six==1.16.0
sympy==1.12
# torch==2.2.1+cu121
# torchaudio==2.2.1+cu121
# torchvision==0.17.1+cu121
tqdm==4.66.5
typing_extensions==4.9.0
tzdata==2024.2
ultralytics==8.3.0
ultralytics-thop==2.0.8
urllib3==2.2.3
onnx>=1.12.0
onnxslim==0.1.34
onnxruntime