diff --git a/.gitignore b/.gitignore index 301f61d..7536cdc 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -p.py \ No newline at end of file +p.py +__pycache__ \ No newline at end of file diff --git a/web_p/blender/assets/assets.json b/web_p/blender/assets/assets.json index fb9ae8f..f661cbb 100644 --- a/web_p/blender/assets/assets.json +++ b/web_p/blender/assets/assets.json @@ -1,6 +1,6 @@ { "assets": [ - { "name": "fork", "mesh": "./mesh/fork.stl", "image": "./images/bear_holder_220425.png" }, + { "name": "bear_holder", "mesh": "./mesh/fork.stl", "image": "./images/bear_holder_220425.png" }, { "name": "bear_holder1", "mesh": "./mesh/fork.stl", "image": "./images/bear_holder_220425.png" } ] } diff --git a/web_p/rbs_train.py b/web_p/rbs_train.py new file mode 100644 index 0000000..2a78759 --- /dev/null +++ b/web_p/rbs_train.py @@ -0,0 +1,29 @@ +""" + rbs_train + Общая задача: web-service pipeline + Реализуемая функция: обучение нейросетевой модели по заданному BOP-датасету + + python3 $PYTHON_EDUCATION --path /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \ + --name test1234 --datasetName 32123213 + + 27.04.2024 @shalenikol release 0.1 +""" +import argparse +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="Path for dataset") + parser.add_argument("--name", required=True, help="String with result weights name") + parser.add_argument("--datasetName", required=True, help="String with dataset name") + parser.add_argument("--outpath", default="weights", help="Output path for weights") + parser.add_argument("--type", default="ObjectDetection", help="Type of implementation") + parser.add_argument("--epoch", default=3, type=int, help="How many training epochs") + parser.add_argument('--pretrain', action="store_true", help="Use pretraining") + args = parser.parse_args() + + if args.type == "ObjectDetection": + train_YoloV8(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain) + else: + train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain) diff --git a/web_p/renderBOPdataset.py b/web_p/renderBOPdataset.py index f3d8b15..56b2ffd 100755 --- a/web_p/renderBOPdataset.py +++ b/web_p/renderBOPdataset.py @@ -5,7 +5,7 @@ import blenderproc as bproc Реализуемая функция: создание датасета в формате BOP с заданными параметрами рандомизации Используется модуль blenderproc - 19.04.2024 @shalenikol release 0.1 + 02.05.2024 @shalenikol release 0.1 """ import numpy as np import argparse @@ -15,7 +15,6 @@ import shutil import json VHACD_PATH = "blenderproc_resources/vhacd" -# DIR_BOP = "bop_data" DIR_MODELS = "models" FILE_LOG_SCENE = "res.txt" FILE_RBS_INFO = "rbs_info.json" @@ -26,15 +25,13 @@ Not_Categories_Name = True # наименование категории в COCO def _get_path_model(name_model: str) -> str: # TODO on name_model find path for mesh (model.fbx) # local_path/assets/mesh/ - return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_model+".fbx") - # , d: dict - # return d["model"] + loc = os.path.dirname(os.path.dirname(rnd_par.output_dir)) + return os.path.join(loc, "assets/mesh/"+name_model+".fbx") def _get_path_object(name_obj: str) -> str: # TODO on name_obj find path for scene object (object.fbx) - return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_obj+".fbx") - # , d: dict - # return d["path"] + loc = os.path.dirname(os.path.dirname(rnd_par.output_dir)) + return os.path.join(loc, "assets/mesh/"+name_obj+".fbx") def convert2relative(height, width, bbox): """ @@ -198,8 +195,6 @@ def render() -> int: t = [obj.get_bound_box(local_coords=True).tolist() for obj in all_meshs if obj.get_name() == objn] rec["cuboid"] = t[0] data.append(rec) - # ff = os.path.join(args.obj_path, rnd_par.models.filenames[i]) # путь к исходному файлу - # shutil.copy2(ff, models_dir) shutil.copy2(rnd_par.models.filenames[i], models_dir) f = (os.path.splitext(rnd_par.models.filenames[i]))[0] + ".mtl" # файл материала if os.path.isfile(f): @@ -260,21 +255,20 @@ def _get_models(par, data) -> int: return 0 # no models # загрузим объекты - par.models.names = [] #list(map(lambda x: x["name"], data)) # obj_names - par.models.filenames = [] #list(map(lambda x: x["model"], data)) #obj_filenames + par.models.names = [] # obj_names + par.models.filenames = [] # obj_filenames i = 1 for f in data: nam = f par.models.names.append(nam) ff = _get_path_model(nam) - # ff = f["model"] # путь к файлу объекта par.models.filenames.append(ff) if not os.path.isfile(ff): print(f"Error: no such file '{ff}'") return -1 obj = bproc.loader.load_obj(ff) all_meshs += obj - obj[0].set_cp("category_id", i) #f["id"]) # начиная с 1 + obj[0].set_cp("category_id", i) # начиная с 1 i += 1 return par.models.n_item @@ -293,7 +287,7 @@ def _get_scene(par, data) -> int: par.scene.objs = [] par.scene.collision_objects = [] for f in objs: - ff = _get_path_object(f["name"]) # f["path"] + ff = _get_path_object(f["name"]) if not os.path.isfile(ff): print(f"Error: no such file '{ff}'") return -1 @@ -303,7 +297,7 @@ def _get_scene(par, data) -> int: if len(coll) > 0: obj[0].enable_rigidbody(False, collision_shape=coll) par.scene.collision_objects += obj - par.scene.objs += obj #bproc.loader.load_blend(args.scene, data_blocks=["objects"]) + par.scene.objs += obj if not par.scene.collision_objects: print("Collision objects not found in the scene") @@ -353,7 +347,6 @@ if __name__ == "__main__": rnd_par.loc_range_high = models_randomization["loc_range_high"] if not os.path.isdir(rnd_par.output_dir): - # os.mkdir(rnd_par.output_dir) print(f"Error: invalid path '{rnd_par.output_dir}'") exit(-3) diff --git a/web_p/train_Dope.py b/web_p/train_Dope.py new file mode 100644 index 0000000..aaea238 --- /dev/null +++ b/web_p/train_Dope.py @@ -0,0 +1,30 @@ +""" + train_Dope + Общая задача: оценка позиции объекта (Pose estimation) + Реализуемая функция: обучение нейросетевой модели DOPE по заданному BOP-датасету + + python3 $PYTHON_EDUCATION --path /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \ + --name test1234 --datasetName 32123213 + + 25.04.2024 @shalenikol release 0.1 +""" +import os + +def train_Dope_i(path:str, wname:str, dname:str, outpath:str, epochs:int): + results = f"torchrun --nproc_per_node=1 train.py --local_rank 0 --data {os.path.join(path,dname)} --object fork" \ + + f" -e {epochs} --batchsize 16 --exts jpg --imagesize 640 --pretrained" \ + + " --net_path /home/shalenikol/fork_work/dope_training/output/weights_2996/net_epoch_47.pth" + print(results) + +import argparse + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--path", required=True, help="Path for dataset") + parser.add_argument("--name", required=True, help="String with result weights name") + parser.add_argument("--datasetName", required=True, help="String with dataset name") + parser.add_argument("--outpath", default="weights", help="Output path for weights") + parser.add_argument("--epoch", default=3, help="How many training epochs") + args = parser.parse_args() + + train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch) diff --git a/web_p/train_Yolo.py b/web_p/train_Yolo.py new file mode 100644 index 0000000..7482603 --- /dev/null +++ b/web_p/train_Yolo.py @@ -0,0 +1,182 @@ +""" + train_Yolo + Общая задача: обнаружение объекта (Object detection) + Реализуемая функция: обучение нейросетевой модели YoloV8 по заданному BOP-датасету + + python3 $PYTHON_TRAIN --path /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296/datasets \ + --name test123 --datasetName ds213 --outpath /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296/weights + + 27.04.2024 @shalenikol release 0.1 +""" +import os +import shutil +import json +import yaml + +from ultralytics import YOLO +# from ultralytics.utils.metrics import DetMetrics + +FILE_BASEMODEL = "yolov8n.pt" +FILE_RBS_INFO = "rbs_info.json" +FILE_RBS_TRAIN = "rbs_train.yaml" +FILE_GT_COCO = "scene_gt_coco.json" +FILE_L_TRAIN = "i_train.txt" +FILE_L_VAL = "i_val.txt" +FILE_TRAIN_RES = "weights/last.pt" +DIR_ROOT_DS = "datasets" +DIR_COCO_DS = "rbs_coco" +DIR_RGB_DS = "images" +DIR_LABELS_DS = "labels" + +SZ_SERIES = 5 # number of train images per validation images + +nn_image = 0 +f1 = f2 = None + +def convert2relative(height, width, bbox): + """ YOLO format use relative coordinates for annotation """ + x, y, w, h = bbox + x += w/2 + y += h/2 + return x/width, y/height, w/width, h/height + +def gt_parse(path: str, out_dir: str): + global nn_image, f1, f2 + with open(os.path.join(path, FILE_GT_COCO), "r") as fh: + coco_data = json.load(fh) + + for img in coco_data["images"]: + rgb_file = os.path.join(path, img["file_name"]) + if os.path.isfile(rgb_file): + ext = os.path.splitext(rgb_file)[1] # only ext + f = f"{nn_image:06}" + out_img = os.path.join(out_dir, DIR_RGB_DS, f + ext) + shutil.copy2(rgb_file, out_img) + + # заполним файлы с метками bbox + img_id = img["id"] + with open(os.path.join(out_dir, DIR_LABELS_DS, f + ".txt"), "w") as fh: + for i in coco_data["annotations"]: + if i["image_id"] == img_id: + cat_id = i["category_id"] + if cat_id < 999: + bbox = i["bbox"] + im_h = i["height"] + im_w = i["width"] + rel = convert2relative(im_h,im_w,bbox) + # формат: + fh.write(f"{cat_id-1} {rel[0]} {rel[1]} {rel[2]} {rel[3]}\n") # category from 0 + + nn_image += 1 + line = os.path.join("./", DIR_RGB_DS, f + ext) + "\n" + if nn_image % SZ_SERIES == 0: + f2.write(line) + else: + f1.write(line) + +def explore(path: str, res_dir: str): + if not os.path.isdir(path): + return + folders = [ + os.path.join(path, o) + for o in os.listdir(path) + if os.path.isdir(os.path.join(path, o)) + ] + for path_entry in folders: + if os.path.isfile(os.path.join(path_entry,FILE_GT_COCO)): + gt_parse(path_entry, res_dir) + else: + explore(path_entry, res_dir) + +def BOP2Yolo_dataset(dpath: str, out_dir: str, lname: list) -> str: + """ Convert BOP-dataset to YOLO format for train """ + cfg_yaml = os.path.join(out_dir, FILE_RBS_TRAIN) + p = os.path.join(out_dir, DIR_ROOT_DS, DIR_COCO_DS) + cfg_data = {"path": p, "train": FILE_L_TRAIN, "val": FILE_L_VAL} + cfg_data["names"] = {i:x for i,x in enumerate(lname)} + with open(cfg_yaml, "w") as fh: + yaml.dump(cfg_data, fh) + + res_dir = os.path.join(out_dir, DIR_ROOT_DS) + if not os.path.isdir(res_dir): + os.mkdir(res_dir) + + res_dir = os.path.join(res_dir, DIR_COCO_DS) + if not os.path.isdir(res_dir): + os.mkdir(res_dir) + + p = os.path.join(res_dir, DIR_RGB_DS) + if not os.path.isdir(p): + os.mkdir(p) + p = os.path.join(res_dir, DIR_LABELS_DS) + if not os.path.isdir(p): + os.mkdir(p) + + global f1, f2 + f1 = open(os.path.join(res_dir, FILE_L_TRAIN), "w") + f2 = open(os.path.join(res_dir, FILE_L_VAL), "w") + explore(dpath, res_dir) + f1.close() + f2.close() + + return out_dir + +def train_YoloV8(path:str, wname:str, dname:str, outpath:str, epochs:int, pretrain: bool): + """ Main procedure for train YOLOv8 model """ + if not os.path.isdir(outpath): + print(f"Invalid output path '{outpath}'") + exit(-1) + out_dir = os.path.join(outpath, wname) + + if pretrain: + # продолжить обучение + if not os.path.isdir(out_dir): + print(f"No dir '{out_dir}'") + exit(-2) + dpath = out_dir + model_path = os.path.join(dpath, wname + ".pt") + else: + # обучение сначала + if not os.path.isdir(out_dir): + os.mkdir(out_dir) + + ds_path = os.path.join(path, dname) + rbs_info = os.path.join(ds_path, FILE_RBS_INFO) + if not os.path.isfile(rbs_info): + print(f"{rbs_info} : no dataset description file") + exit(-3) + + with open(rbs_info, "r") as fh: + y = json.load(fh) + # список имён объектов + list_name = list(map(lambda x: x["name"], y)) + + dpath = BOP2Yolo_dataset(ds_path, out_dir, list_name) + if len(dpath) == 0: + print(f"Error in convert dataset '{ds_path}' to '{outpath}'") + exit(-4) + model_path = os.path.join(dpath, FILE_BASEMODEL) + + model = YOLO(model_path) + results = model.train(data=os.path.join(dpath, FILE_RBS_TRAIN), epochs=epochs, project=out_dir) + wf = os.path.join(results.save_dir, FILE_TRAIN_RES) + if not os.path.isfile(wf): + print(f"Error in train: no result file '{wf}'") + exit(-5) + + shutil.copy2(wf, os.path.join(dpath, wname + ".pt")) + shutil.rmtree(results.save_dir) + # print(f"\n ********\n{wf}") + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--path", required=True, help="Path for dataset") + parser.add_argument("--name", required=True, help="String with result weights name") + parser.add_argument("--datasetName", required=True, help="String with dataset name") + parser.add_argument("--outpath", default="weights", help="Output path for weights") + parser.add_argument("--epoch", default=3, type=int, help="How many training epochs") + parser.add_argument('--pretrain', action="store_true", help="Use pretraining") + args = parser.parse_args() + + train_YoloV8(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain)