107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
from ultralytics import YOLO
|
||
import os
|
||
import torch
|
||
from pathlib import Path
|
||
|
||
def detect_images():
|
||
# ======================= 配置区 =======================
|
||
# 模型配置
|
||
model_config = {
|
||
'model_path': 'yolo11x.pt', # 本地模型路径
|
||
'download_url':'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11x.pt' # 如果没有模型文件可在此处下载URL
|
||
}
|
||
|
||
# 路径配置
|
||
path_config = {
|
||
'input_folder': 'input',
|
||
'output_folder': 'output',
|
||
'auto_create_dir': True # 自动创建输出目录
|
||
}
|
||
|
||
# 推理参数
|
||
predict_config = {
|
||
'conf_thres': 0.25, # 置信度阈值
|
||
'iou_thres': 0.45, # IoU阈值
|
||
'imgsz': 640, # 输入分辨率
|
||
'line_width': 2, # 检测框线宽
|
||
'device': 'cuda:0' if torch.cuda.is_available() else 'cpu' # 自动选择设备
|
||
}
|
||
# ====================== 配置结束 ======================
|
||
|
||
try:
|
||
# 验证输入目录
|
||
if not Path(path_config['input_folder']).exists():
|
||
raise FileNotFoundError(f"输入目录不存在: {path_config['input_folder']}")
|
||
|
||
# 自动创建输出目录
|
||
if path_config['auto_create_dir']:
|
||
Path(path_config['output_folder']).mkdir(parents=True, exist_ok=True)
|
||
|
||
# 加载模型（带异常捕获）
|
||
if not Path(model_config['model_path']).exists():
|
||
if model_config['download_url']:
|
||
print("开始下载模型...")
|
||
YOLO(model_config['download_url']).download(model_config['model_path'])
|
||
else:
|
||
raise FileNotFoundError(f"模型文件不存在: {model_config['model_path']}")
|
||
|
||
# 初始化模型
|
||
model = YOLO(model_config['model_path']).to(predict_config['device'])
|
||
print(f"✅ 模型加载成功 | 设备: {predict_config['device'].upper()}")
|
||
|
||
# 执行推理
|
||
results = model.predict(
|
||
source=path_config['input_folder'],
|
||
project=path_config['output_folder'],
|
||
name="exp",
|
||
save=True,
|
||
conf=predict_config['conf_thres'],
|
||
iou=predict_config['iou_thres'],
|
||
imgsz=predict_config['imgsz'],
|
||
line_width=predict_config['line_width'],
|
||
show_labels=True,
|
||
show_conf=True,
|
||
classes=[2] # 只检测car类别
|
||
)
|
||
|
||
# 只保留每张图片中面积最大的car框
|
||
output_txt = os.path.join(path_config['output_folder'], 'car_boxes.txt')
|
||
with open(output_txt, 'w') as f:
|
||
for result in results:
|
||
boxes = result.boxes.xyxy.cpu().numpy() # [N, 4]
|
||
path = result.path
|
||
img_name = os.path.basename(path)
|
||
if len(boxes) > 0:
|
||
# 计算面积
|
||
areas = [(box[2] - box[0]) * (box[3] - box[1]) for box in boxes]
|
||
max_idx = areas.index(max(areas))
|
||
box = boxes[max_idx]
|
||
x1, y1, x2, y2 = box
|
||
points = [
|
||
(x1, y1), # 左上
|
||
(x2, y2), # 右下
|
||
]
|
||
f.write(f"{img_name} " + " ".join([f"{int(x)} {int(y)}" for x, y in points]) + "\n")
|
||
print(f"已保存car框坐标到: {output_txt}")
|
||
|
||
# 统计信息
|
||
success_count = len(results)
|
||
save_dir = Path(results[0].save_dir) if success_count > 0 else None
|
||
print(f"\n🔍 推理完成 | 处理图片: {success_count} 张")
|
||
print(f"📁 结果目录: {save_dir.resolve() if save_dir else '无'}")
|
||
|
||
# 显示首张结果（可选）
|
||
if success_count > 0:
|
||
results[0].show()
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ 发生错误: {str(e)}")
|
||
print("问题排查建议：")
|
||
print("1. 检查模型文件路径是否正确")
|
||
print("2. 确认图片目录包含支持的格式（jpg/png等）")
|
||
print("3. 查看CUDA是否可用（如需GPU加速）")
|
||
print("4. 确保输出目录有写入权限")
|
||
|
||
if __name__ == "__main__":
|
||
detect_images()
|