alexander
This commit is contained in:
IDONTSUDO
parent
c49beb8218
commit
e0a6cd0af1
6 changed files with 254 additions and 19 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -1 +1,2 @@
|
|||
p.py
|
||||
p.py
|
||||
__pycache__
|
||||
|
|
@ -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" }
|
||||
]
|
||||
}
|
||||
|
|
|
|||
29
web_p/rbs_train.py
Normal file
29
web_p/rbs_train.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
30
web_p/train_Dope.py
Normal file
30
web_p/train_Dope.py
Normal file
|
|
@ -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)
|
||||
182
web_p/train_Yolo.py
Normal file
182
web_p/train_Yolo.py
Normal file
|
|
@ -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)
|
||||
# формат: <target> <x-center> <y-center> <width> <height>
|
||||
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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue