64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
"""
|
|
rbs_train2
|
|
Общая задача: web-service pipeline
|
|
Реализуемая функция: обучение нейросетевой модели по заданному BOP-датасету
|
|
|
|
python3 $PYTHON_EDUCATION --path /home/user/webservice/server/build/public/process/proc/inst_proc \
|
|
--form /home/user/webservice/server/build/public/process/proc/inst_proc/form.json
|
|
|
|
28.01.2025 @shalenikol release 0.1
|
|
17.02.2025 @shalenikol release 0.2 addon_dir
|
|
"""
|
|
import argparse
|
|
import os
|
|
import json
|
|
from train_Yolo import train_YoloV8
|
|
from train_Dope import train_Dope_i
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--path", required=True, help="Output path for weights")
|
|
parser.add_argument("--form", required=True, help="Json-file with training parameters")
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.isdir(args.path):
|
|
print(f"Invalid output path '{args.path}'")
|
|
exit(-1)
|
|
wname = os.path.basename(args.path)
|
|
outpath = os.path.dirname(args.path)
|
|
|
|
if not os.path.isfile(args.form):
|
|
print(f"Error: no such file '{args.form}'")
|
|
exit(-2)
|
|
with open(args.form, "r") as f:
|
|
j_data = f.read()
|
|
try:
|
|
cfg = json.loads(j_data)
|
|
except json.JSONDecodeError as e:
|
|
print(f"JSon error: {e}")
|
|
exit(-3)
|
|
|
|
cfg = cfg["output"] # edited params
|
|
dataset_params = cfg["process"]["selectProcess"]["value"]
|
|
dataset_type = dataset_params["type"]
|
|
if dataset_type != "BOP_DATASET":
|
|
print(f"Error: Invalid dataset type '{dataset_type}'")
|
|
exit(-4)
|
|
dataset_name = dataset_params["instanceName"]
|
|
dataset_path = dataset_params["path"]
|
|
dataset_path = dataset_path.replace("//", "/") # !!! TODO !!! Nikita
|
|
|
|
epoch = cfg["n_epoch"]
|
|
pretrain = (cfg["pretrain"] == "True") #False
|
|
ttype = cfg["typeWeight"] #"ObjectDetection"
|
|
|
|
addon_dir = ""
|
|
if "addon" in cfg:
|
|
addon = cfg["addon"].strip()
|
|
if addon and os.path.isdir(addon):
|
|
addon_dir = addon
|
|
|
|
if ttype == "ObjectDetection":
|
|
train_YoloV8(dataset_path, wname, dataset_name, outpath, epoch, pretrain, addon_dir)
|
|
else:
|
|
train_Dope_i(dataset_path, wname, dataset_name, outpath, epoch, pretrain)
|