#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2023/5/4 20:27:02
@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
import laok
from laok.cv2d.io import fix_cv_read
from ultralytics import YOLO
#===============================================================================
r'''支持 yolo训练的脚本
需要 文件夹 root_dir下:
                ---- xxx.pt
                ---- xxx.names
                ---- Image/
'''
#===============================================================================

def _gen_data(root_dir, filename = None):
    root_dir = laok.path_abs(root_dir)
    names_file = laok.path_search_cur_ext(root_dir, ".txt;.names")
    dataset = laok.Dict()
    dataset.path = laok.path_native(root_dir)
    dataset.train = 'Image'
    dataset.val = 'Image'
    dataset.test = ''
    dataset.names = list(laok.file_read_lines(names_file))
    if filename is None:
        filename = laok.path_replace_ext(names_file, ".yaml")
    laok.save_yaml_file(filename, dataset, indent=4, encoding='utf8')
    return filename

def train_det(root_dir, model_file=None, data_file=None, imgsz=(512, 1248), epochs=100, batch=16, **kwargs):
    print(f'root_dir: {root_dir}')

    if model_file is None:
        model_file = laok.path_search_cur_ext(root_dir, ".pt")
    print(f'model_file: {model_file}')

    fix_cv_read()
    with laok.work_dir_scope(root_dir):
        data_file = _gen_data(root_dir, data_file)
        print(f'data_file: {data_file}' )

        model = YOLO(model_file)  # load a pretrained model (recommended for training)
        kwargs['task'] = 'detect'

        kwargs.setdefault('save_txt', True)
        kwargs.setdefault('save_conf', True)
        kwargs.setdefault('save_crop', True)
        kwargs.setdefault('amp', False)
        if laok.is_windows():
            kwargs.setdefault('workers', 0)
        model.train(data=data_file, epochs=epochs, imgsz=imgsz, batch=batch, **kwargs)
        best_model = laok.path_search_under_file(model.trainer.save_dir, 'best.pt')
        model2 = YOLO(best_model)  # load a custom trained
        model2.export(format='onnx', opset=12, dynamic = False, imgsz=imgsz)

if __name__ == '__main__':
    import fire as _fire
    _fire.Fire(train_det)