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（输出路径）')
|