121 lines
3.5 KiB
Python
121 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
为训练图片自动生成标注
|
|
使用 YOLO 模型检测野鸡并生成 YOLO 格式标注
|
|
"""
|
|
|
|
import cv2
|
|
import os
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
|
|
# 配置
|
|
MODEL_PATH = Path(__file__).parent / "yolov8s-world.pt"
|
|
IMAGES_DIR = Path(__file__).parent / "datasets" / "images" / "pheasant"
|
|
OUTPUT_DIR = Path(__file__).parent / "datasets" / "yolo_format"
|
|
|
|
def create_yolo_structure():
|
|
"""创建 YOLO 格式目录结构"""
|
|
(OUTPUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
|
|
(OUTPUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)
|
|
|
|
def convert_to_yolo_format(bbox, img_width, img_height):
|
|
"""将边界框转换为 YOLO 格式 (cx, cy, w, h)"""
|
|
x1, y1, x2, y2 = bbox
|
|
cx = (x1 + x2) / 2 / img_width
|
|
cy = (y1 + y2) / 2 / img_height
|
|
w = (x2 - x1) / img_width
|
|
h = (y2 - y1) / img_height
|
|
return cx, cy, w, h
|
|
|
|
def main():
|
|
# 加载模型
|
|
print(f"加载模型: {MODEL_PATH}")
|
|
model = YOLO(str(MODEL_PATH))
|
|
|
|
# 创建目录结构
|
|
create_yolo_structure()
|
|
|
|
# 获取所有图片
|
|
image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
|
|
image_files = [
|
|
f for f in IMAGES_DIR.iterdir()
|
|
if f.suffix.lower() in image_extensions
|
|
]
|
|
|
|
print(f"找到 {len(image_files)} 张图片")
|
|
|
|
# 处理每张图片
|
|
labeled_count = 0
|
|
for img_path in image_files:
|
|
print(f"处理: {img_path.name}")
|
|
|
|
# 读取图片
|
|
image = cv2.imread(str(img_path))
|
|
if image is None:
|
|
print(f" 跳过: 无法读取 {img_path.name}")
|
|
continue
|
|
|
|
img_height, img_width = image.shape[:2]
|
|
|
|
# 推理
|
|
results = model(image, conf=0.3, iou=0.45)
|
|
|
|
# 收集标注
|
|
labels = []
|
|
for result in results:
|
|
boxes = result.boxes
|
|
if boxes is None or len(boxes) == 0:
|
|
continue
|
|
|
|
for box in boxes:
|
|
class_name = result.names.get(int(box.cls[0]), "")
|
|
|
|
# 只保留 bird 或 pheasant 类别
|
|
if class_name in ["bird", "pheasant"]:
|
|
class_id = 0 # 只有一个类别：野鸡
|
|
bbox = list(map(int, box.xyxy[0].tolist()))
|
|
conf = float(box.conf[0])
|
|
|
|
# 转换为 YOLO 格式
|
|
cx, cy, w, h = convert_to_yolo_format(bbox, img_width, img_height)
|
|
labels.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
|
|
|
# 保存标注文件
|
|
if len(labels) > 0:
|
|
labeled_count += 1
|
|
# 复制图片
|
|
dst_img_path = OUTPUT_DIR / "images" / "train" / img_path.name
|
|
cv2.imwrite(str(dst_img_path), image)
|
|
|
|
# 保存标注
|
|
label_path = OUTPUT_DIR / "labels" / "train" / (img_path.stem + ".txt")
|
|
with open(label_path, "w") as f:
|
|
f.write("\n".join(labels))
|
|
|
|
print(f" ✓ 标注了 {len(labels)} 个目标")
|
|
else:
|
|
print(f" - 未检测到野鸡")
|
|
|
|
print(f"\n完成!")
|
|
print(f" 总图片数: {len(image_files)}")
|
|
print(f" 有效标注: {labeled_count}")
|
|
print(f" 输出目录: {OUTPUT_DIR}")
|
|
|
|
# 创建 data.yaml
|
|
data_yaml = OUTPUT_DIR / "data.yaml"
|
|
with open(data_yaml, "w") as f:
|
|
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
|
path: {OUTPUT_DIR}
|
|
train: images/train
|
|
val: images/train
|
|
|
|
# 类别
|
|
nc: 1
|
|
names: ['pheasant']
|
|
""")
|
|
print(f" 数据配置: {data_yaml}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|