52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
|
import os
|
|||
|
|
from config import BUCKET_NAME, LOCAL_DATASET_DIR
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
from tools import validate_and_update_yaml_fields, convert_xml_to_txt
|
|||
|
|
from minio_tools import download_minio_folder
|
|||
|
|
from dataset_action import update_database_status
|
|||
|
|
import ray
|
|||
|
|
from yolov5.train import really_training
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
@ray.remote(max_calls=1, num_gpus=0.2)
|
|||
|
|
def training(training_id, minio_folder_name, coco_path, pre_model_path, images_path, labels_path, epochs, img_size, batch_size):
|
|||
|
|
|
|||
|
|
## 1、首先判断一下该数据集是否已经下载
|
|||
|
|
if os.path.exists(images_path) and os.path.isdir(images_path):
|
|||
|
|
print(f"数据集已存在。")
|
|||
|
|
else:
|
|||
|
|
print(f"数据集不存在，开始下载。")
|
|||
|
|
if not os.path.exists(LOCAL_DATASET_DIR):
|
|||
|
|
os.makedirs(LOCAL_DATASET_DIR)
|
|||
|
|
is_load = download_minio_folder(BUCKET_NAME, minio_folder_name, LOCAL_DATASET_DIR)
|
|||
|
|
if not is_load:
|
|||
|
|
update_database_status(training_id, "初始化异常")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
## 2、更新yaml文件，将其中的路径替换为实际训练路径
|
|||
|
|
train_path = os.path.join(images_path, "train")
|
|||
|
|
val_path = os.path.join(images_path, "val")
|
|||
|
|
is_val, categories = validate_and_update_yaml_fields(coco_path, train_path, val_path) # 从这里顺便获取类别
|
|||
|
|
if not is_val:
|
|||
|
|
update_database_status(training_id, "初始化异常")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
## 3、检查数据集和图片是否完整（将xml文件转为txt文件）
|
|||
|
|
is_convert = convert_xml_to_txt(categories, labels_path)
|
|||
|
|
if not is_convert:
|
|||
|
|
update_database_status(training_id, "初始化异常")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
## 4、开始训练
|
|||
|
|
update_database_status(training_id, "训练中")
|
|||
|
|
result = really_training(training_id, pre_model_path, coco_path, epochs, img_size, batch_size)
|
|||
|
|
update_database_status(training_id, "训练完成")
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
pass
|