60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
训练 YOLOv8n 模型 - 只识别野鸡（pheasant）
|
||
使用现有的训练数据进行迁移学习
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from ultralytics import YOLO
|
||
|
||
# 配置
|
||
DATA_DIR = Path(__file__).parent / "datasets"
|
||
MODEL_NAME = "yolov8n.pt" # 预训练模型
|
||
OUTPUT_DIR = Path(__file__).parent / "runs"
|
||
|
||
def main():
|
||
# 创建数据集配置文件（只包含野鸡类别）
|
||
data_yaml = DATA_DIR / "data.yaml"
|
||
print("创建数据集配置文件（只训练野鸡类别）...")
|
||
with open(data_yaml, "w") as f:
|
||
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
||
path: {DATA_DIR}
|
||
train: images/pheasant
|
||
val: images/pheasant
|
||
|
||
# 类别
|
||
nc: 1
|
||
names: ['pheasant']
|
||
""")
|
||
|
||
# 加载预训练模型
|
||
print(f"加载预训练模型: {MODEL_NAME}")
|
||
model = YOLO(MODEL_NAME)
|
||
|
||
# 训练模型
|
||
print("开始训练...")
|
||
results = model.train(
|
||
data=str(data_yaml),
|
||
epochs=50,
|
||
imgsz=640,
|
||
batch=16,
|
||
name="observer_yolov8n",
|
||
patience=20,
|
||
save=True,
|
||
plots=True
|
||
)
|
||
|
||
print(f"\n训练完成!")
|
||
print(f"最佳模型保存在: {OUTPUT_DIR / 'observer_yolov8n' / 'weights' / 'best.pt'}")
|
||
|
||
# 导出为 TFLite 格式
|
||
print("\n导出为 TFLite 格式...")
|
||
best_model_path = OUTPUT_DIR / "observer_yolov8n" / "weights" / "best.pt"
|
||
if best_model_path.exists():
|
||
best_model = YOLO(str(best_model_path))
|
||
best_model.export(format="tflite", imgsz=320)
|
||
print(f"TFLite 模型导出完成")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|