86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
from ultralytics import YOLO
|
||
import cv2
|
||
import os
|
||
|
||
def detect_and_save(image_path, output_dir="result_yolov8", save_txt=False):
|
||
"""
|
||
使用 YOLOv8 检测图片中的目标并保存结果
|
||
|
||
参数:
|
||
image_path: 输入图片路径
|
||
output_dir: 结果保存目录
|
||
"""
|
||
# 加载预训练模型 (自动下载 yolov8n.pt 若不存在)
|
||
model = YOLO("./yolov8n.pt")
|
||
|
||
# 进行目标检测
|
||
results = model(image_path)
|
||
|
||
# 创建输出目录
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 处理每个检测结果 (单张图片返回一个结果列表)
|
||
for i, r in enumerate(results):
|
||
# 1. 保存带标注的图片
|
||
annotated_img = r.plot() # 生成带标注框的BGR图像
|
||
img_name = os.path.basename(image_path)
|
||
img_output_path = os.path.join(output_dir, f"annotated_{img_name}.jpg")
|
||
print(f"保存带标注的图片：{img_output_path}")
|
||
cv2.imwrite(img_output_path, annotated_img)
|
||
print(f"✅ 标注图片已保存至: {img_output_path}")
|
||
|
||
if save_txt:
|
||
# 2. 保存检测结果(txt文件，YOLO格式)
|
||
txt_path = os.path.join(output_dir, f"results_{os.path.splitext(img_name)[0]}.txt")
|
||
r.save_txt(txt_path)
|
||
print(f"📝 检测结果已保存至: {txt_path}")
|
||
|
||
# 3. 控制台打印检测信息
|
||
boxes = r.boxes
|
||
print("\n检测到的目标:")
|
||
print(f"{'类别':<10}{'置信度':<15}{'坐标(xywh)':<30}")
|
||
print("-" * 55)
|
||
for box in boxes:
|
||
cls_id = int(box.cls)
|
||
class_name = model.names[cls_id]
|
||
|
||
# 统计类别信息
|
||
class_set.add(class_name)
|
||
|
||
conf = float(box.conf)
|
||
bbox = [int(coord) for coord in box.xywh[0]] # 获取xywh格式坐标
|
||
print(f"{class_name:<10}{conf:.4f}{'':<5}{bbox}")
|
||
print("-" * 55)
|
||
|
||
|
||
def detect_file(path_file, output_dir):
|
||
t = os.listdir(path_file)
|
||
for i in t:
|
||
detect_and_save(os.path.join(path_file, i), os.path.join(output_dir, path_file.split("/")[-1]))
|
||
|
||
|
||
def detect_file_list(path_file, output_dir):
|
||
t = os.listdir(path_file)
|
||
for i in t:
|
||
detect_file(os.path.join(path_file, i), output_dir)
|
||
|
||
|
||
class_set = set()
|
||
|
||
if __name__ == "__main__":
|
||
import os
|
||
|
||
path_file = "data_image"
|
||
output_dir = "result_yolov8"
|
||
|
||
detect_file_list(path_file, output_dir)
|
||
|
||
print(class_set)
|
||
|
||
# # 输入图片路径 (可替换为你的图片路径)
|
||
# input_image = "data_image/1e4c75b76e531606e2adc491a8f09ae8/frame_000150.jpg" # 确保图片存在
|
||
# output_dir = "result_yolov8"
|
||
|
||
|
||
# # 执行检测并保存结果
|
||
# detect_and_save(input_image, output_dir) |