107 lines
4.4 KiB
Python
107 lines
4.4 KiB
Python
|
#静态检测，图片检测
|
|||
|
|
|
|||
|
|
from ultralytics import YOLO # ultralytics是yolo模型的实现库
|
|||
|
|
import cv2 # cv2是opencv库，用于图像处理
|
|||
|
|
import os # os用于文件和目录操作
|
|||
|
|
|
|||
|
|
# 替换为实际的模型路径
|
|||
|
|
model_path = r'\detect\train\weights\best.pt' # 加载yolo模型
|
|||
|
|
model = YOLO(model_path) # model是加载后的yolo模型实例
|
|||
|
|
|
|||
|
|
def infer_and_draw(image_path, output_folder):
|
|||
|
|
# 检查 image_path 是否是文件夹路径
|
|||
|
|
if os.path.isdir(image_path):
|
|||
|
|
# 遍历文件夹中的所有图像文件
|
|||
|
|
for filename in os.listdir(image_path):
|
|||
|
|
file_path = os.path.join(image_path, filename)
|
|||
|
|
if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
|
|||
|
|
process_image(file_path, output_folder)
|
|||
|
|
else:
|
|||
|
|
# 处理单张图像
|
|||
|
|
process_image(image_path, output_folder)
|
|||
|
|
|
|||
|
|
def process_image(image_path, output_folder):
|
|||
|
|
# 使用模型进行推理
|
|||
|
|
results = model(image_path)
|
|||
|
|
original_img = cv2.imread(image_path) # 读取原始图像
|
|||
|
|
|
|||
|
|
if results:
|
|||
|
|
for result in results:
|
|||
|
|
if result.boxes is not None and len(result.boxes.xyxy) > 0:
|
|||
|
|
boxes = result.boxes.xyxy.cpu().numpy() # 获取检测框
|
|||
|
|
classes = result.boxes.cls.cpu().numpy() # 获取类别
|
|||
|
|
confidences = result.boxes.conf.cpu().numpy() # 获取置信度
|
|||
|
|
|
|||
|
|
# 确保输出文件夹存在
|
|||
|
|
if not os.path.exists(output_folder):
|
|||
|
|
os.makedirs(output_folder)
|
|||
|
|
|
|||
|
|
# 存储所有检测框的标签和位置
|
|||
|
|
detected_boxes = []
|
|||
|
|
|
|||
|
|
for i, box in enumerate(boxes):
|
|||
|
|
x1, y1, x2, y2 = map(int, box[:4])
|
|||
|
|
class_id = int(classes[i]) if len(classes) > i else 0
|
|||
|
|
# 标签映射（假设0表示人头类别）
|
|||
|
|
label = "head" if class_id == 0 else "unknown"
|
|||
|
|
conf = confidences[i] if len(confidences) > i else 0.0
|
|||
|
|
|
|||
|
|
detected_boxes.append((x1, y1, x2, y2, label, conf))
|
|||
|
|
|
|||
|
|
# 避免重叠框的问题，通过过滤和合并框
|
|||
|
|
filtered_boxes = filter_and_merge_boxes(detected_boxes)
|
|||
|
|
|
|||
|
|
# 生成最终的标注图像
|
|||
|
|
img_annotated = original_img.copy()
|
|||
|
|
for (x1, y1, x2, y2, label, conf) in filtered_boxes:
|
|||
|
|
cv2.rectangle(img_annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|||
|
|
cv2.putText(img_annotated, f'{label} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|||
|
|
|
|||
|
|
# 保存标注后的图像
|
|||
|
|
output_image_path = os.path.join(output_folder, os.path.basename(image_path))
|
|||
|
|
cv2.imwrite(output_image_path, img_annotated)
|
|||
|
|
print(f"Saved annotated image to: {output_image_path}")
|
|||
|
|
|
|||
|
|
# 显示图像（可选）
|
|||
|
|
cv2.imshow('Detected Image', img_annotated)
|
|||
|
|
cv2.waitKey(0)
|
|||
|
|
cv2.destroyAllWindows()
|
|||
|
|
else:
|
|||
|
|
print("No detections found.")
|
|||
|
|
# 保存原始图像
|
|||
|
|
output_image_path = os.path.join(output_folder, os.path.basename(image_path))
|
|||
|
|
cv2.imwrite(output_image_path, original_img)
|
|||
|
|
print(f"Saved original image to: {output_image_path}")
|
|||
|
|
|
|||
|
|
def filter_and_merge_boxes(boxes):
|
|||
|
|
"""
|
|||
|
|
简单地合并重叠框的逻辑。可以根据需要调整合并算法。
|
|||
|
|
"""
|
|||
|
|
filtered_boxes = []
|
|||
|
|
threshold = 0.5 # 重叠阈值
|
|||
|
|
|
|||
|
|
def iou(box1, box2):
|
|||
|
|
x1, y1, x2, y2 = box1
|
|||
|
|
x1_, y1_, x2_, y2_ = box2
|
|||
|
|
ix1, iy1 = max(x1, x1_), max(y1, y1_)
|
|||
|
|
ix2, iy2 = min(x2, x2_), min(y2, y2_)
|
|||
|
|
iw = max(ix2 - ix1 + 1, 0)
|
|||
|
|
ih = max(iy2 - iy1 + 1, 0)
|
|||
|
|
inter = iw * ih
|
|||
|
|
ua = (x2 - x1 + 1) * (y2 - y1 + 1) + (x2_ - x1_ + 1) * (y2_ - y1_) - inter
|
|||
|
|
return inter / ua
|
|||
|
|
|
|||
|
|
for i, box1 in enumerate(boxes):
|
|||
|
|
keep = True
|
|||
|
|
for j, box2 in enumerate(filtered_boxes):
|
|||
|
|
if iou(box1[:4], box2[:4]) > threshold:
|
|||
|
|
keep = False
|
|||
|
|
break
|
|||
|
|
if keep:
|
|||
|
|
filtered_boxes.append(box1)
|
|||
|
|
|
|||
|
|
return filtered_boxes
|
|||
|
|
|
|||
|
|
# 使用实际图像路径或文件夹路径进行推理，并指定输出文件夹
|
|||
|
|
infer_and_draw(r'\assets\xx.jpg(图片路径)', r'ultralytics\output（输出路径）')
|