131 lines
5.2 KiB
Python
131 lines
5.2 KiB
Python
import os
|
||
from ultralytics import YOLO
|
||
import cv2
|
||
import random
|
||
|
||
def generate_label_colors(labels):
|
||
"""为每个标签生成固定颜色（红色，黄色，蓝色，绿色）"""
|
||
colors = {
|
||
"red": (255, 0, 0),
|
||
"yellow": (255, 255, 0),
|
||
"blue": (0, 0, 255),
|
||
"green": (0, 255, 0),
|
||
}
|
||
# 为标签分配颜色，确保颜色在红、黄、蓝、绿之间循环
|
||
label_colors = {}
|
||
color_list = list(colors.values())
|
||
for i, label in enumerate(labels):
|
||
print(f"generate_label_colors, label: {label}")
|
||
label_colors[label] = color_list[i % len(color_list)] # 循环分配颜色
|
||
|
||
return label_colors
|
||
|
||
def intelligent_annotation_with_filter(model_path, input_folder, output_folder, target_labels, conf_threshold=0.3):
|
||
"""
|
||
使用 YOLOv8 模型对本地文件夹图像进行智能标注，并按指定标签过滤输出。
|
||
|
||
:param model_path: YOLOv8 模型文件路径
|
||
:param input_folder: 输入文件夹路径，包含待标注的图像
|
||
:param output_folder: 输出文件夹路径，存储标注后的图像
|
||
:param target_labels: 用户指定需要标注的目标类别列表
|
||
:param conf_threshold: 置信度阈值，默认 0.5
|
||
"""
|
||
# 加载 YOLO 模型
|
||
model = YOLO(model_path)
|
||
|
||
# 获取所有可能的标签
|
||
all_labels = model.names.values()
|
||
print("所有可能的标签：", all_labels)
|
||
label_colors = generate_label_colors(all_labels)
|
||
print("所有可能的标签对应的颜色：", label_colors)
|
||
|
||
# 保存标注结果的字典
|
||
annotation_results = {}
|
||
|
||
# 创建输出文件夹（如果不存在）
|
||
os.makedirs(output_folder, exist_ok=True)
|
||
|
||
# 遍历输入文件夹中的所有图像
|
||
for file_name in os.listdir(input_folder):
|
||
input_path = os.path.join(input_folder, file_name)
|
||
|
||
# 检查文件是否是图像格式
|
||
if not (file_name.endswith('.jpg') or file_name.endswith('.png') or file_name.endswith('.jpeg')):
|
||
continue
|
||
|
||
# 模型预测
|
||
results = model(input_path, task="detect", conf=conf_threshold)
|
||
|
||
image_annotations = [] # 当前图像的标注信息
|
||
|
||
# 读取原始图像
|
||
image = cv2.imread(input_path)
|
||
|
||
# 获取图像尺寸
|
||
image_height, image_width = image.shape[:2]
|
||
|
||
# 标注的标志，检查是否有符合的目标
|
||
has_annotation = False
|
||
|
||
# 遍历预测结果
|
||
for result in results:
|
||
for box in result.boxes:
|
||
# 提取边界框信息和分类标签
|
||
x1, y1, x2, y2 = map(int, box.xyxy[0]) # 坐标转换为整数
|
||
label = model.names[int(box.cls[0])] # 类别名称
|
||
print(f"label:" + label)
|
||
conf = box.conf[0] # 置信度
|
||
|
||
# 检查是否为目标标签
|
||
if label in target_labels:
|
||
color = label_colors[label] # 获取对应标签的颜色
|
||
has_annotation = True
|
||
|
||
# 绘制边界框
|
||
# color = (0, 255, 0) # 绿色
|
||
# 随机生成颜色
|
||
# color = tuple(random.randint(0, 255) for _ in range(3)) # 生成随机颜色
|
||
cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 在边界框上方绘制标签和置信度
|
||
label_text = f"{label} {conf:.2f}"
|
||
print(f"label_text:" + label_text)
|
||
cv2.putText(image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 1)
|
||
|
||
# 转换为归一化坐标并保存
|
||
normalized_x = x1 / image_width
|
||
normalized_y = y1 / image_height
|
||
normalized_width = (x2 - x1) / image_width
|
||
normalized_height = (y2 - y1) / image_height
|
||
|
||
annotation = {
|
||
"label": label,
|
||
"x": normalized_x * 100,
|
||
"y": normalized_y * 100,
|
||
"width": normalized_width * 100,
|
||
"height": normalized_height * 100,
|
||
}
|
||
image_annotations.append(annotation)
|
||
|
||
# 保存标注信息到结果字典
|
||
annotation_results[input_path] = image_annotations if image_annotations else []
|
||
|
||
# 如果存在符合的标注目标，保存标注后的图像
|
||
if has_annotation:
|
||
output_path = os.path.join(output_folder, file_name)
|
||
cv2.imwrite(output_path, image)
|
||
|
||
print(f"智能标注完成！标注后的图像保存在：{output_folder}")
|
||
return annotation_results
|
||
|
||
model_path = "./model/sishu_thin_24_12.onnx" # 你的YOLOv8模型权重路径
|
||
input_folder = "./dataset" # 输入文件夹路径
|
||
output_folder = "./output" # 输出文件夹路径
|
||
target_labels = ["tiaojuan", "zhujiesi", "yulinwen"] # 用户指定的目标标签
|
||
annotation_results = intelligent_annotation_with_filter(model_path, input_folder, output_folder, target_labels)
|
||
print(f"标注结果：annotation_results：{annotation_results}")
|
||
|
||
print(f"标注结果：./dataset/1024072305_0724123933.jpg annotation_results：{annotation_results['./dataset/1024072305_0724123933.jpg']}")
|
||
|
||
|