mirror of
https://github.com/m5stack/UNITV_TRAIN.git
synced 2026-05-20 11:05:22 -07:00
61 lines
2.5 KiB
Python
61 lines
2.5 KiB
Python
#!python3
|
|
|
|
import argparse, os, sys, shutil
|
|
|
|
|
|
def main():
|
|
supported_types = ["classifier", "detector"]
|
|
curr_dir = os.path.abspath(os.path.dirname(__file__))
|
|
parser = argparse.ArgumentParser(description="train model", usage='''
|
|
python3 main.py -z "datasets zip file" init
|
|
then
|
|
python3 main.py -z "datasets zip file" train
|
|
or python3 main.py -d "datasets directory" train
|
|
''')
|
|
parser.add_argument("-t", "--type", type=str, help="train type, classifier or detector", choices=supported_types, default="classifier")
|
|
parser.add_argument("-z", "--zip", type=str, help="datasets zip file path", default="")
|
|
parser.add_argument("-d", "--datasets", type=str, help="datasets directory", default="")
|
|
parser.add_argument("-c", "--config", type=str, help="config file", default=os.path.join(curr_dir, "instance", "config.py"))
|
|
parser.add_argument("-o", "--out", type=str, help="out directory", default=os.path.join(curr_dir, "out"))
|
|
parser.add_argument("cmd", help="command", choices=["train", "init"])
|
|
args = parser.parse_args()
|
|
# init
|
|
dst_config_path = os.path.join(curr_dir, "instance", "config.py")
|
|
if args.cmd == "init":
|
|
instance_dir = os.path.join(curr_dir, "instance")
|
|
if not os.path.exists(instance_dir):
|
|
os.makedirs(instance_dir)
|
|
copy_config = True
|
|
if os.path.exists(dst_config_path):
|
|
print("[WARNING] instance/config.py already exists, sure to rewrite it? [yes/no]")
|
|
ensure = input()
|
|
if ensure != "yes":
|
|
copy_config = False
|
|
if copy_config:
|
|
shutil.copyfile(os.path.join(curr_dir, "train", "config_template.py"), dst_config_path)
|
|
print("init done, please edit instance/config.py")
|
|
return 0
|
|
if not os.path.exists(dst_config_path):
|
|
print("python3 train.py init first")
|
|
return -1
|
|
|
|
from train import Train, TrainType
|
|
is_zip = os.path.exists(args.zip)
|
|
if not is_zip and not os.path.exists(args.datasets):
|
|
print("[ERROR] arg -d or -z is essential")
|
|
return -1
|
|
if args.type == "classifier":
|
|
train_task = Train(TrainType.CLASSIFIER, args.zip, args.datasets, args.out)
|
|
elif args.type == "detector":
|
|
train_task = Train(TrainType.DETECTOR, args.zip, args.datasets, args.out)
|
|
else:
|
|
print("[ERROR] train type not support only support: {}".format(", ".join(supported_types)))
|
|
train_task.train()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|
|
|
|
|