800 lines
28 KiB
Python
800 lines
28 KiB
Python
|
# import asyncio
|
|||
|
|
# import os.path
|
|||
|
|
# import shutil
|
|||
|
|
# import sys
|
|||
|
|
# import threading
|
|||
|
|
# import time
|
|||
|
|
# from pathlib import Path
|
|||
|
|
#
|
|||
|
|
# from hachoir.parser.image.iptc import datasets
|
|||
|
|
#
|
|||
|
|
# from middleware.minio_util import downFullPathFile
|
|||
|
|
# from middleware.query_model import ModelConfigDAO
|
|||
|
|
# import yaml
|
|||
|
|
# import multiprocessing
|
|||
|
|
# import torch
|
|||
|
|
# from ultralytics import YOLO
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# async def download_train(task_id: str, bz_training_task_id: int, pt_name: str):
|
|||
|
|
#
|
|||
|
|
# DB_CONFIG = {
|
|||
|
|
# "dbname": "smart_dev_123",
|
|||
|
|
# "user": "postgres",
|
|||
|
|
# "password": "root",
|
|||
|
|
# "host": "8.137.54.85",
|
|||
|
|
# "port": "5060"
|
|||
|
|
# }
|
|||
|
|
#
|
|||
|
|
# # 创建DAO实例
|
|||
|
|
# dao = ModelConfigDAO(DB_CONFIG)
|
|||
|
|
# time_ns=time.time_ns()
|
|||
|
|
# output_root=f"dataset-{time_ns}"
|
|||
|
|
# if not os.path.exists(output_root):
|
|||
|
|
# os.mkdir(output_root)
|
|||
|
|
# list_labels = dao.get_labels(bz_training_task_id)
|
|||
|
|
# list_datasets = dao.get_datasets(bz_training_task_id)
|
|||
|
|
# label_yaml_list = dao.get_label_yaml(bz_training_task_id)
|
|||
|
|
#
|
|||
|
|
# # 定义数据结构（字典）
|
|||
|
|
# uavid_config = {
|
|||
|
|
#
|
|||
|
|
# "path": "", # 替换为你的绝对路径
|
|||
|
|
# "train": "images/train", # 训练集路径
|
|||
|
|
# "val": "images/val", # 验证集路径
|
|||
|
|
# "test": "images/test", # 测试集路径（可选）
|
|||
|
|
# "names": {}
|
|||
|
|
# }
|
|||
|
|
# uavid_config["path"]=os.path.abspath(output_root)
|
|||
|
|
# for i,item in enumerate(label_yaml_list):
|
|||
|
|
# item.id_order=i
|
|||
|
|
# uavid_config["names"][f"{i}"]=item.e_name
|
|||
|
|
# # 生成 YAML 文件
|
|||
|
|
# data_yaml="data.yaml"
|
|||
|
|
# with open(data_yaml, "w", encoding="utf-8") as f:
|
|||
|
|
# yaml.dump(
|
|||
|
|
# uavid_config,
|
|||
|
|
# f,
|
|||
|
|
# default_flow_style=False, # 禁用紧凑格式（保持多行）
|
|||
|
|
# allow_unicode=True, # 允许 Unicode 字符
|
|||
|
|
# sort_keys=False # 保持键的顺序
|
|||
|
|
# )
|
|||
|
|
# file_name = os.path.basename(data_yaml)
|
|||
|
|
# des_path = os.path.join(output_root, file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_yaml, output_root)
|
|||
|
|
# print(f"output_rootoutput_rootoutput_rootoutput_root {os.path.abspath(output_root)}")
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# print("YAML 文件已生成：uavid_config.yaml")
|
|||
|
|
#
|
|||
|
|
# invalid_indices = []
|
|||
|
|
#
|
|||
|
|
# for index, pic in enumerate(list_datasets):
|
|||
|
|
# if pic.resource_original_path: # 图像路径有效
|
|||
|
|
# download_path = downFullPathFile(pic.resource_original_path)
|
|||
|
|
# if download_path: # 下载成功
|
|||
|
|
# pic.local_path = download_path
|
|||
|
|
#
|
|||
|
|
# pic.label_name = Path(download_path).stem # 截取图片名称，用作标签
|
|||
|
|
# else:
|
|||
|
|
# invalid_indices.append(index) #存储不符合条件的索引，准备删除
|
|||
|
|
# else:
|
|||
|
|
# invalid_indices.append(index)#存储不符合条件的索引，准备删除
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# # 从后往前删除（避免删除时索引错乱）,删除不符合条件的list_datasets
|
|||
|
|
#
|
|||
|
|
# for idx in sorted(invalid_indices, reverse=True):
|
|||
|
|
# del list_datasets[idx]
|
|||
|
|
#
|
|||
|
|
# for data_pic in list_datasets: #整理完整的图像与标签集的对应关系
|
|||
|
|
# for label in list_labels:
|
|||
|
|
# if data_pic.id == label.id:
|
|||
|
|
# for item in label_yaml_list:
|
|||
|
|
# if label.label_ids==item.id:
|
|||
|
|
# data_pic.label_content=data_pic.label_content+item.id_order+" "+label.annotation_data+ '\n'
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# for data_pic in list_datasets:
|
|||
|
|
# label_txt = f"{data_pic.label_name}.txt"
|
|||
|
|
# with open(label_txt, 'w', encoding='utf-8') as f:
|
|||
|
|
# f.write(data_pic.label_content)
|
|||
|
|
# data_pic.label_txt_path=os.path.abspath(label_txt)
|
|||
|
|
# # 移动文件，制作数据集
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# dataset_dirs = {
|
|||
|
|
# "images": Path(output_root) / "images",
|
|||
|
|
# "labels": Path(output_root) / "labels"
|
|||
|
|
# }
|
|||
|
|
# for ds_dir in dataset_dirs.values():
|
|||
|
|
# (ds_dir / "val").mkdir(parents=True, exist_ok=True)
|
|||
|
|
# (ds_dir / "train").mkdir(parents=True, exist_ok=True)
|
|||
|
|
# (ds_dir / "test").mkdir(parents=True, exist_ok=True)
|
|||
|
|
#
|
|||
|
|
# count_pic=0
|
|||
|
|
# for data_pic in list_datasets:
|
|||
|
|
# count_pic=count_pic+1
|
|||
|
|
# if count_pic%10<8:
|
|||
|
|
# images_train_path=dataset_dirs["images"]
|
|||
|
|
# image_dir=os.path.join(images_train_path,"train")
|
|||
|
|
# file_name=os.path.basename(data_pic.local_path)
|
|||
|
|
# des_path=os.path.join(image_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.local_path, image_dir)
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# labels_train_path=dataset_dirs["labels"]
|
|||
|
|
# label_dir=os.path.join(labels_train_path,"train")
|
|||
|
|
#
|
|||
|
|
# file_name=os.path.basename(data_pic.label_txt_path)
|
|||
|
|
# des_path=os.path.join(label_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.label_txt_path, label_dir)
|
|||
|
|
# if count_pic%10==8:
|
|||
|
|
# images_val_path=dataset_dirs["images"]
|
|||
|
|
# image_dir=os.path.join(images_val_path,"val")
|
|||
|
|
# file_name=os.path.basename(data_pic.local_path)
|
|||
|
|
# des_path=os.path.join(image_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.local_path, image_dir)
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# labels_val_path=dataset_dirs["labels"]
|
|||
|
|
# label_dir=os.path.join(labels_val_path,"val")
|
|||
|
|
# file_name=os.path.basename(data_pic.label_txt_path)
|
|||
|
|
# des_path=os.path.join(label_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.label_txt_path, label_dir)
|
|||
|
|
#
|
|||
|
|
# if count_pic%10==9:
|
|||
|
|
# images_test_path=dataset_dirs["images"]
|
|||
|
|
# image_dir=os.path.join(images_test_path,"test")
|
|||
|
|
# file_name=os.path.basename(data_pic.local_path)
|
|||
|
|
# des_path=os.path.join(image_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.local_path, image_dir)
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# labels_test_path=dataset_dirs["labels"]
|
|||
|
|
# label_dir=os.path.join(labels_test_path,"test")
|
|||
|
|
# file_name=os.path.basename(data_pic.label_txt_path)
|
|||
|
|
# des_path=os.path.join(label_dir,file_name)
|
|||
|
|
# if os.path.exists(des_path):
|
|||
|
|
# os.remove(des_path)
|
|||
|
|
# shutil.move(data_pic.label_txt_path, label_dir)
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# custom_config = {
|
|||
|
|
# "epochs": 50, # 快速测试用
|
|||
|
|
# "batch_size": 4,
|
|||
|
|
# }
|
|||
|
|
#
|
|||
|
|
# # 启动后台训练
|
|||
|
|
# pid = await run_background_training(
|
|||
|
|
# dataset_dir=output_root,
|
|||
|
|
# weight_name=pt_name,
|
|||
|
|
# config_overrides=custom_config
|
|||
|
|
# )
|
|||
|
|
#
|
|||
|
|
# print(f"pid--{pid}")
|
|||
|
|
# dao.insert_train_pid(task_id,train_pid=pid)
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# def train_model(dataset_dir,weight_name="best_segmentation_model.pt", config_overrides=None):
|
|||
|
|
# """
|
|||
|
|
# 训练模型并保存权重
|
|||
|
|
# :param weight_name: 自定义权重文件名（如 "uavid_seg_v1.pt"）
|
|||
|
|
# :param config_overrides: 覆盖默认配置的字典（可选）
|
|||
|
|
# """
|
|||
|
|
# # 合并配置（允许通过参数覆盖默认配置）
|
|||
|
|
#
|
|||
|
|
# # 默认配置（可通过函数参数覆盖）
|
|||
|
|
# DEFAULT_CONFIG = {
|
|||
|
|
# "model": "pt/yolo11s-seg.pt",
|
|||
|
|
# "pretrained": True,
|
|||
|
|
# "data": os.path.join(dataset_dir, "data.yaml"), # 关键修改：指向 data.yaml
|
|||
|
|
# "project": "UAVid_Segmentation",
|
|||
|
|
# "name": "v1.5_official",
|
|||
|
|
# "epochs": 1000,
|
|||
|
|
# "batch_size": 8,
|
|||
|
|
# "img_size": 640,
|
|||
|
|
# "workers": 4,
|
|||
|
|
# "optimizer": "SGD",
|
|||
|
|
# "lr0": 0.01,
|
|||
|
|
# "lrf": 0.01,
|
|||
|
|
# "momentum": 0.9,
|
|||
|
|
# "weight_decay": 0.0005,
|
|||
|
|
# "augment": True,
|
|||
|
|
# "hyp": {
|
|||
|
|
# "mosaic": 0.5,
|
|||
|
|
# "copy_paste": 0.2,
|
|||
|
|
# "mixup": 0.15,
|
|||
|
|
# },
|
|||
|
|
# }
|
|||
|
|
#
|
|||
|
|
# config = DEFAULT_CONFIG.copy()
|
|||
|
|
# if config_overrides:
|
|||
|
|
# config.update(config_overrides)
|
|||
|
|
#
|
|||
|
|
# # 初始化模型
|
|||
|
|
# model = YOLO(config["model"])
|
|||
|
|
#
|
|||
|
|
# # 开始训练
|
|||
|
|
# results = model.train(
|
|||
|
|
# data=config["data"],
|
|||
|
|
# project=config["project"],
|
|||
|
|
# name=config["name"],
|
|||
|
|
# epochs=config["epochs"],
|
|||
|
|
# batch=config["batch_size"],
|
|||
|
|
# imgsz=config["img_size"],
|
|||
|
|
# workers=config["workers"],
|
|||
|
|
# optimizer=config["optimizer"],
|
|||
|
|
# lr0=config["lr0"],
|
|||
|
|
# lrf=config["lrf"],
|
|||
|
|
# momentum=config["momentum"],
|
|||
|
|
# weight_decay=config["weight_decay"],
|
|||
|
|
# augment=config["augment"],
|
|||
|
|
# device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
|
|||
|
|
# )
|
|||
|
|
#
|
|||
|
|
# # 验证模型
|
|||
|
|
# metrics = model.val()
|
|||
|
|
# print(f"Validation mAP: {metrics.box_map:.2f} (box), {metrics.seg_map:.2f} (mask)")
|
|||
|
|
#
|
|||
|
|
# # 保存最佳模型（使用自定义名称）
|
|||
|
|
# best_model = results.best_model
|
|||
|
|
# torch.save(best_model, weight_name)
|
|||
|
|
# print(f"Best model saved to: {weight_name}")
|
|||
|
|
#
|
|||
|
|
#
|
|||
|
|
# # def run_background_training(output_root: str, weight_name="best_segmentation_model.pt", config_overrides=None):
|
|||
|
|
# # """使用 spawn 上下文启动进程"""
|
|||
|
|
# # ctx = multiprocessing.get_context('spawn')
|
|||
|
|
# # process = ctx.Process(
|
|||
|
|
# # target=train_model,
|
|||
|
|
# # args=(output_root, weight_name, config_overrides),
|
|||
|
|
# # daemon=False
|
|||
|
|
# # )
|
|||
|
|
# # process.start()
|
|||
|
|
# # return process.pid
|
|||
|
|
#
|
|||
|
|
# import asyncio
|
|||
|
|
#
|
|||
|
|
# async def run_background_training(dataset_dir, weight_name, config_overrides=None):
|
|||
|
|
# """异步启动训练进程"""
|
|||
|
|
# process = await asyncio.create_subprocess_exec(
|
|||
|
|
# sys.executable,
|
|||
|
|
# "train_worker.py",
|
|||
|
|
# "--dataset", dataset_dir,
|
|||
|
|
# "--weight", weight_name,
|
|||
|
|
# "--epochs", str(config_overrides.get("epochs", 50)),
|
|||
|
|
# "--batch", str(config_overrides.get("batch", 4)),
|
|||
|
|
# stdout=asyncio.subprocess.PIPE,
|
|||
|
|
# stderr=asyncio.subprocess.PIPE,
|
|||
|
|
# )
|
|||
|
|
# return process.pid
|
|||
|
|
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import os.path
|
|||
|
|
import shutil
|
|||
|
|
import sys
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
import subprocess
|
|||
|
|
import json
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
from middleware.minio_util import downFullPathFile
|
|||
|
|
from middleware.query_model import ModelConfigDAO
|
|||
|
|
import yaml
|
|||
|
|
import torch
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def download_train(task_id: str, bz_training_task_id: int, pt_name: str):
|
|||
|
|
"""
|
|||
|
|
下载训练数据并启动训练
|
|||
|
|
这个函数负责准备数据，然后使用线程+subprocess创建独立进程执行训练
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
current_pid = os.getpid()
|
|||
|
|
print(f"Starting download and training for task {task_id} in process {current_pid}")
|
|||
|
|
|
|||
|
|
DB_CONFIG = {
|
|||
|
|
"dbname": "smart_dev_123",
|
|||
|
|
"user": "postgres",
|
|||
|
|
"password": "root",
|
|||
|
|
"host": "8.137.54.85",
|
|||
|
|
"port": "5060"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 创建DAO实例
|
|||
|
|
dao = ModelConfigDAO(DB_CONFIG)
|
|||
|
|
time_ns = time.time_ns()
|
|||
|
|
|
|||
|
|
output_root = f"dataset/dataset-{time_ns}"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if not os.path.exists(output_root):
|
|||
|
|
os.mkdir(output_root)
|
|||
|
|
print(f"Created output directory: {output_root}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Failed to create output directory: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 获取标签和数据集信息
|
|||
|
|
list_labels = dao.get_labels(bz_training_task_id)
|
|||
|
|
list_datasets = dao.get_datasets(bz_training_task_id)
|
|||
|
|
label_yaml_list = dao.get_label_yaml(bz_training_task_id)
|
|||
|
|
|
|||
|
|
print(
|
|||
|
|
f"Retrieved {len(list_labels)} labels, {len(list_datasets)} datasets, {len(label_yaml_list)} label configs")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Failed to retrieve data from database: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 定义数据结构（字典）
|
|||
|
|
uavid_config = {
|
|||
|
|
"path": "", # 替换为你的绝对路径
|
|||
|
|
"train": "images/train", # 训练集路径
|
|||
|
|
"val": "images/val", # 验证集路径
|
|||
|
|
"test": "images/test", # 测试集路径（可选）
|
|||
|
|
"names": {}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
uavid_config["path"] = os.path.abspath(output_root)
|
|||
|
|
for i, item in enumerate(label_yaml_list):
|
|||
|
|
item.id_order = i
|
|||
|
|
uavid_config["names"][f"{i}"] = item.e_name
|
|||
|
|
|
|||
|
|
# 生成 YAML 文件
|
|||
|
|
data_yaml = "data.yaml"
|
|||
|
|
with open(data_yaml, "w", encoding="utf-8") as f:
|
|||
|
|
yaml.dump(
|
|||
|
|
uavid_config,
|
|||
|
|
f,
|
|||
|
|
default_flow_style=False, # 禁用紧凑格式（保持多行）
|
|||
|
|
allow_unicode=True, # 允许 Unicode 字符
|
|||
|
|
sort_keys=False # 保持键的顺序
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
file_name = os.path.basename(data_yaml)
|
|||
|
|
des_path = os.path.join(output_root, file_name)
|
|||
|
|
if os.path.exists(des_path):
|
|||
|
|
os.remove(des_path)
|
|||
|
|
shutil.move(data_yaml, output_root)
|
|||
|
|
print(f"Generated YAML config at: {os.path.abspath(output_root)}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Failed to generate YAML config: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 下载数据集
|
|||
|
|
invalid_indices = []
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
for index, pic in enumerate(list_datasets):
|
|||
|
|
if hasattr(pic, 'resource_original_path') and pic.resource_original_path: # 图像路径有效
|
|||
|
|
try:
|
|||
|
|
download_path = downFullPathFile(pic.resource_original_path)
|
|||
|
|
if download_path: # 下载成功
|
|||
|
|
pic.local_path = download_path
|
|||
|
|
pic.label_name = Path(download_path).stem # 截取图片名称，用作标签
|
|||
|
|
print(f"Downloaded file: {download_path}")
|
|||
|
|
else:
|
|||
|
|
invalid_indices.append(index) # 存储不符合条件的索引，准备删除
|
|||
|
|
print(f"Failed to download file: {pic.resource_original_path}")
|
|||
|
|
except Exception as e:
|
|||
|
|
invalid_indices.append(index)
|
|||
|
|
print(f"Error downloading file {pic.resource_original_path}: {e}")
|
|||
|
|
else:
|
|||
|
|
invalid_indices.append(index) # 存储不符合条件的索引，准备删除
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error processing datasets: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 从后往前删除（避免删除时索引错乱）,删除不符合条件的list_datasets
|
|||
|
|
try:
|
|||
|
|
for idx in sorted(invalid_indices, reverse=True):
|
|||
|
|
del list_datasets[idx]
|
|||
|
|
|
|||
|
|
print(f"Filtered datasets: {len(list_datasets)} valid items remaining")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error filtering datasets: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 整理标签内容
|
|||
|
|
try:
|
|||
|
|
for data_pic in list_datasets: # 整理完整的图像与标签集的对应关系
|
|||
|
|
for label in list_labels:
|
|||
|
|
if hasattr(data_pic, 'id') and hasattr(label, 'id') and data_pic.id == label.id:
|
|||
|
|
for item in label_yaml_list:
|
|||
|
|
if hasattr(label, 'label_ids') and hasattr(item, 'id') and label.label_ids == item.id:
|
|||
|
|
# 假设label有annotation_data属性
|
|||
|
|
annotation = getattr(label, 'annotation_data', '')
|
|||
|
|
current_content = getattr(data_pic, 'label_content', '')
|
|||
|
|
data_pic.label_content = f"{current_content}{item.id_order} {annotation}\n"
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error organizing labels: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 创建标签文件
|
|||
|
|
try:
|
|||
|
|
for data_pic in list_datasets:
|
|||
|
|
if hasattr(data_pic, 'label_name'):
|
|||
|
|
label_txt = f"{data_pic.label_name}.txt"
|
|||
|
|
with open(label_txt, 'w', encoding='utf-8') as f:
|
|||
|
|
f.write(getattr(data_pic, 'label_content', ''))
|
|||
|
|
data_pic.label_txt_path = os.path.abspath(label_txt)
|
|||
|
|
print(f"Created label file: {label_txt}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error creating label files: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 移动文件，制作数据集
|
|||
|
|
try:
|
|||
|
|
dataset_dirs = {
|
|||
|
|
"images": Path(output_root) / "images",
|
|||
|
|
"labels": Path(output_root) / "labels"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
for ds_dir in dataset_dirs.values():
|
|||
|
|
(ds_dir / "val").mkdir(parents=True, exist_ok=True)
|
|||
|
|
(ds_dir / "train").mkdir(parents=True, exist_ok=True)
|
|||
|
|
(ds_dir / "test").mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
print("Created dataset directory structure")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error creating dataset directories: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 分配数据集到训练、验证、测试集
|
|||
|
|
try:
|
|||
|
|
count_pic = 0
|
|||
|
|
for data_pic in list_datasets:
|
|||
|
|
count_pic += 1
|
|||
|
|
|
|||
|
|
# 80% 训练集, 10% 验证集, 10% 测试集
|
|||
|
|
if count_pic % 10 < 8:
|
|||
|
|
split = "train"
|
|||
|
|
elif count_pic % 10 == 8:
|
|||
|
|
split = "val"
|
|||
|
|
else: # count_pic % 10 == 9
|
|||
|
|
split = "test"
|
|||
|
|
|
|||
|
|
# 移动图像文件
|
|||
|
|
if hasattr(data_pic, 'local_path') and os.path.exists(data_pic.local_path):
|
|||
|
|
images_path = dataset_dirs["images"]
|
|||
|
|
image_dir = os.path.join(images_path, split)
|
|||
|
|
file_name = os.path.basename(data_pic.local_path)
|
|||
|
|
des_path = os.path.join(image_dir, file_name)
|
|||
|
|
|
|||
|
|
if os.path.exists(des_path):
|
|||
|
|
os.remove(des_path)
|
|||
|
|
shutil.move(data_pic.local_path, image_dir)
|
|||
|
|
|
|||
|
|
# 移动标签文件
|
|||
|
|
if hasattr(data_pic, 'label_txt_path') and os.path.exists(data_pic.label_txt_path):
|
|||
|
|
labels_path = dataset_dirs["labels"]
|
|||
|
|
label_dir = os.path.join(labels_path, split)
|
|||
|
|
file_name = os.path.basename(data_pic.label_txt_path)
|
|||
|
|
des_path = os.path.join(label_dir, file_name)
|
|||
|
|
|
|||
|
|
if os.path.exists(des_path):
|
|||
|
|
os.remove(des_path)
|
|||
|
|
shutil.move(data_pic.label_txt_path, label_dir)
|
|||
|
|
|
|||
|
|
print(f"Organized {count_pic} files into dataset splits")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error organizing dataset splits: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 训练配置
|
|||
|
|
custom_config = {
|
|||
|
|
"epochs": 50, # 快速测试用
|
|||
|
|
"batch_size": 4,
|
|||
|
|
"workers": 0, # 禁用多进程数据加载
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 保存训练配置到文件
|
|||
|
|
config_file = f"train_config_{task_id}.json"
|
|||
|
|
with open(config_file, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump({
|
|||
|
|
'dataset_dir': output_root,
|
|||
|
|
'pt_name': pt_name,
|
|||
|
|
'config_overrides': custom_config,
|
|||
|
|
'db_config': DB_CONFIG,
|
|||
|
|
'task_id': task_id
|
|||
|
|
}, f)
|
|||
|
|
|
|||
|
|
print(f"Training data preparation completed for task {task_id}")
|
|||
|
|
|
|||
|
|
# 在Windows上使用线程+subprocess创建训练进程
|
|||
|
|
# 避免使用asyncio.create_subprocess_exec
|
|||
|
|
loop = asyncio.get_event_loop()
|
|||
|
|
training_pid = await loop.run_in_executor(
|
|||
|
|
None, # 使用默认的线程池
|
|||
|
|
start_training_process,
|
|||
|
|
config_file
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if training_pid:
|
|||
|
|
print(f"pid--{training_pid}")
|
|||
|
|
dao.insert_train_pid(task_id, train_pid=training_pid)
|
|||
|
|
return training_pid
|
|||
|
|
else:
|
|||
|
|
raise Exception("Failed to start training process")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Training failed for task {task_id}: {e}", exc_info=True)
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
def start_training_process(config_file: str) -> int:
|
|||
|
|
"""
|
|||
|
|
在独立线程中启动训练进程
|
|||
|
|
使用subprocess.Popen创建训练进程
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 创建训练脚本内容
|
|||
|
|
train_script = '''
|
|||
|
|
import sys
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import torch
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
class MockModelConfigDAO:
|
|||
|
|
def __init__(self, db_config):
|
|||
|
|
self.db_config = db_config
|
|||
|
|
|
|||
|
|
def insert_train_pid(self, task_id, train_pid):
|
|||
|
|
print(f"Inserted training PID {train_pid} for task {task_id}")
|
|||
|
|
|
|||
|
|
def train_model(dataset_dir, weight_name="best_segmentation_model.pt", config_overrides=None):
|
|||
|
|
"""
|
|||
|
|
训练模型并保存权重
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
current_pid = os.getpid()
|
|||
|
|
print(f"Starting model training in process {current_pid} with dataset: {dataset_dir}")
|
|||
|
|
|
|||
|
|
# 默认配置（可通过参数覆盖）
|
|||
|
|
DEFAULT_CONFIG = {
|
|||
|
|
"model": "pt/yolo11s-seg.pt",
|
|||
|
|
"pretrained": True,
|
|||
|
|
"data": os.path.join(dataset_dir, "data.yaml"),
|
|||
|
|
"project": "UAVid_Segmentation",
|
|||
|
|
"name": "v1.5_official",
|
|||
|
|
"epochs": 1000,
|
|||
|
|
"batch_size": 8,
|
|||
|
|
"img_size": 640,
|
|||
|
|
"workers": 0, # 禁用多进程数据加载
|
|||
|
|
"optimizer": "SGD",
|
|||
|
|
"lr0": 0.01,
|
|||
|
|
"lrf": 0.01,
|
|||
|
|
"momentum": 0.9,
|
|||
|
|
"weight_decay": 0.0005,
|
|||
|
|
"augment": True,
|
|||
|
|
"hyp": {
|
|||
|
|
"mosaic": 0.5,
|
|||
|
|
"copy_paste": 0.2,
|
|||
|
|
"mixup": 0.15,
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
config = DEFAULT_CONFIG.copy()
|
|||
|
|
if config_overrides:
|
|||
|
|
config.update(config_overrides)
|
|||
|
|
|
|||
|
|
print(f"Training config: {config}")
|
|||
|
|
|
|||
|
|
# 检查数据配置文件
|
|||
|
|
data_path = config["data"]
|
|||
|
|
if not os.path.exists(data_path):
|
|||
|
|
raise FileNotFoundError(f"Data configuration file not found: {data_path}")
|
|||
|
|
|
|||
|
|
# 初始化模型
|
|||
|
|
model = YOLO(config["model"])
|
|||
|
|
print(f"Model initialized with: {config["model"]}")
|
|||
|
|
|
|||
|
|
# 开始训练
|
|||
|
|
results = model.train(
|
|||
|
|
data=config["data"],
|
|||
|
|
project=config["project"],
|
|||
|
|
name=config["name"],
|
|||
|
|
epochs=config["epochs"],
|
|||
|
|
batch=config["batch_size"],
|
|||
|
|
imgsz=config["img_size"],
|
|||
|
|
workers=config["workers"],
|
|||
|
|
optimizer=config["optimizer"],
|
|||
|
|
lr0=config["lr0"],
|
|||
|
|
lrf=config["lrf"],
|
|||
|
|
momentum=config["momentum"],
|
|||
|
|
weight_decay=config["weight_decay"],
|
|||
|
|
augment=config["augment"],
|
|||
|
|
device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
|
|||
|
|
)
|
|||
|
|
print(f"Training completed successfully in process {current_pid}")
|
|||
|
|
|
|||
|
|
# 验证模型
|
|||
|
|
metrics = model.val()
|
|||
|
|
print(f"Validation mAP: {metrics.box:.2f} (box), {metrics.seg:.2f} (mask)")
|
|||
|
|
|
|||
|
|
# 保存最佳模型
|
|||
|
|
try:
|
|||
|
|
if hasattr(results, 'best') and results.best:
|
|||
|
|
best_model_path = results.best
|
|||
|
|
if os.path.exists(best_model_path):
|
|||
|
|
import shutil
|
|||
|
|
shutil.copy2(best_model_path, weight_name)
|
|||
|
|
print(f"Best model saved to: {os.path.abspath(weight_name)}")
|
|||
|
|
else:
|
|||
|
|
torch.save(model.state_dict(), weight_name)
|
|||
|
|
print(f"Best model path not found, saved state dict to: {weight_name}")
|
|||
|
|
else:
|
|||
|
|
torch.save(model.state_dict(), weight_name)
|
|||
|
|
print(f"Saved model state dict to: {weight_name}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Warning: Failed to save best model: {e}")
|
|||
|
|
torch.save(model.state_dict(), weight_name)
|
|||
|
|
print(f"Fallback: Saved model state dict to: {weight_name}")
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Model training failed in process {os.getpid()}: {e}", exc_info=True)
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
if len(sys.argv) != 2:
|
|||
|
|
print("Usage: python -c '<script>' <config_file>")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
config_file = sys.argv[1]
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with open(config_file, 'r', encoding='utf-8') as f:
|
|||
|
|
config = json.load(f)
|
|||
|
|
|
|||
|
|
# 提取配置
|
|||
|
|
dataset_dir = config['dataset_dir']
|
|||
|
|
pt_name = config['pt_name']
|
|||
|
|
config_overrides = config['config_overrides']
|
|||
|
|
db_config = config['db_config']
|
|||
|
|
task_id = config['task_id']
|
|||
|
|
|
|||
|
|
# 获取当前进程ID
|
|||
|
|
pid = os.getpid()
|
|||
|
|
print(f"Training process started for task {task_id} with PID {pid}")
|
|||
|
|
|
|||
|
|
# 记录PID到数据库
|
|||
|
|
try:
|
|||
|
|
from middleware.query_model import ModelConfigDAO
|
|||
|
|
dao = ModelConfigDAO(db_config)
|
|||
|
|
except ImportError:
|
|||
|
|
dao = MockModelConfigDAO(db_config)
|
|||
|
|
|
|||
|
|
dao.insert_train_pid(task_id, train_pid=pid)
|
|||
|
|
|
|||
|
|
# 执行训练
|
|||
|
|
success = train_model(dataset_dir, pt_name, config_overrides)
|
|||
|
|
|
|||
|
|
if success:
|
|||
|
|
print(f"Training completed successfully for task {task_id}")
|
|||
|
|
sys.exit(0)
|
|||
|
|
else:
|
|||
|
|
print(f"Training failed for task {task_id}")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Training error: {e}", file=sys.stderr)
|
|||
|
|
import traceback
|
|||
|
|
traceback.print_exc()
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|
|||
|
|
'''
|
|||
|
|
|
|||
|
|
# 保存训练脚本
|
|||
|
|
script_path = f"train_worker_{os.path.basename(config_file).split('_')[2].split('.')[0]}.py"
|
|||
|
|
with open(script_path, 'w', encoding='utf-8') as f:
|
|||
|
|
f.write(train_script)
|
|||
|
|
|
|||
|
|
# 使用subprocess.Popen创建训练进程
|
|||
|
|
# 在Windows上使用shell=True可以解决一些路径问题
|
|||
|
|
process = subprocess.Popen([
|
|||
|
|
sys.executable, script_path, config_file
|
|||
|
|
], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=False)
|
|||
|
|
|
|||
|
|
print(f"Started training process with PID {process.pid}")
|
|||
|
|
|
|||
|
|
# 启动线程来处理输出
|
|||
|
|
threading.Thread(target=handle_process_output, args=(process,), daemon=True).start()
|
|||
|
|
|
|||
|
|
return process.pid
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Failed to start training process: {e}", exc_info=True)
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
def handle_process_output(process: subprocess.Popen):
|
|||
|
|
"""
|
|||
|
|
处理子进程的输出
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 分别读取stdout和stderr
|
|||
|
|
def read_stream(stream, stream_name):
|
|||
|
|
while True:
|
|||
|
|
line = stream.readline()
|
|||
|
|
if not line:
|
|||
|
|
break
|
|||
|
|
line = line.strip()
|
|||
|
|
print(f"[{stream_name}] {line}")
|
|||
|
|
|
|||
|
|
# 启动线程读取stdout和stderr
|
|||
|
|
stdout_thread = threading.Thread(target=read_stream, args=(process.stdout, 'STDOUT'))
|
|||
|
|
stderr_thread = threading.Thread(target=read_stream, args=(process.stderr, 'STDERR'))
|
|||
|
|
|
|||
|
|
stdout_thread.start()
|
|||
|
|
stderr_thread.start()
|
|||
|
|
|
|||
|
|
# 等待进程完成
|
|||
|
|
stdout_thread.join()
|
|||
|
|
stderr_thread.join()
|
|||
|
|
|
|||
|
|
# 获取返回码
|
|||
|
|
return_code = process.wait()
|
|||
|
|
print(f"Training process completed with return code: {return_code}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Error handling process output: {e}", exc_info=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
"""
|
|||
|
|
直接运行时的入口点
|
|||
|
|
用于测试
|
|||
|
|
"""
|
|||
|
|
import sys
|
|||
|
|
|
|||
|
|
if len(sys.argv) < 4:
|
|||
|
|
print(f"Usage: {sys.argv[0]} <task_id> <train_task_id> <pt_name>")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
task_id = sys.argv[1]
|
|||
|
|
train_task_id = int(sys.argv[2])
|
|||
|
|
pt_name = sys.argv[3]
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 创建事件循环
|
|||
|
|
loop = asyncio.get_event_loop()
|
|||
|
|
pid = loop.run_until_complete(download_train(task_id, train_task_id, pt_name))
|
|||
|
|
print(f"Training started in process {pid}")
|
|||
|
|
|
|||
|
|
# 保持事件循环运行
|
|||
|
|
try:
|
|||
|
|
loop.run_forever()
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print("Received keyboard interrupt. Exiting...")
|
|||
|
|
finally:
|
|||
|
|
loop.close()
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Training failed: {e}", exc_info=True)
|
|||
|
|
sys.exit(1)
|
|||
|
|
|