update scripts in web_p (dataset, train)
This commit is contained in:
parent
7d3b8ff0cb
commit
c85784f3dc
3 changed files with 256 additions and 29 deletions
|
|
@ -7,6 +7,8 @@
|
|||
--name test123 --datasetName ds213 --outpath /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296/weights
|
||||
|
||||
27.04.2024 @shalenikol release 0.1
|
||||
20.11.2024 @shalenikol release 0.2 parser.add_argument("--addon", default="", help="Folder with add-on for dataset")
|
||||
20.02.2025 @shalenikol release 0.2.1 add_on_dataset : fix
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
|
|
@ -14,9 +16,13 @@ import json
|
|||
import yaml
|
||||
|
||||
from ultralytics import YOLO
|
||||
# from ultralytics import settings
|
||||
# from ultralytics.utils.metrics import DetMetrics
|
||||
# import torch
|
||||
# import torch.profiler
|
||||
# import torch.utils.data
|
||||
|
||||
FILE_BASEMODEL = "yolov8n.pt"
|
||||
FILE_BASEMODEL = "yolov8s.pt" #"yolov8n.pt"
|
||||
FILE_RBS_INFO = "rbs_info.json"
|
||||
FILE_RBS_TRAIN = "rbs_train.yaml"
|
||||
FILE_GT_COCO = "scene_gt_coco.json"
|
||||
|
|
@ -27,6 +33,7 @@ DIR_ROOT_DS = "datasets"
|
|||
DIR_COCO_DS = "rbs_coco"
|
||||
DIR_RGB_DS = "images"
|
||||
DIR_LABELS_DS = "labels"
|
||||
LABELS_EXT = ".txt"
|
||||
|
||||
SZ_SERIES = 15 # number of train images per validation images
|
||||
|
||||
|
|
@ -40,6 +47,50 @@ def convert2relative(height, width, bbox):
|
|||
y += h/2
|
||||
return x/width, y/height, w/width, h/height
|
||||
|
||||
def add_on_dataset(source_dir, target_dir) -> dict:
|
||||
global nn_image, f1, f2
|
||||
# Получаем список файлов в исходной директории
|
||||
files = sorted(os.listdir(source_dir))
|
||||
|
||||
# Словарь для отслеживания порядковых номеров для каждого имени файла
|
||||
file_nn = {}
|
||||
|
||||
for file in files:
|
||||
if os.path.isdir(os.path.join(source_dir, file)):
|
||||
continue
|
||||
# Получаем имя файла и его расширение
|
||||
file_name, file_extension = os.path.splitext(file)
|
||||
|
||||
# Запоминаем порядковый номер для данного имени файла
|
||||
if file_name in file_nn:
|
||||
nn = file_nn[file_name]
|
||||
else: # new file name
|
||||
nn = nn_image # текущий номер
|
||||
file_nn[file_name] = nn_image
|
||||
nn_image += 1
|
||||
|
||||
# Создаем новое имя файла
|
||||
new_file_name = f"{nn:06}{file_extension}"
|
||||
if file_extension == LABELS_EXT:
|
||||
new_file_path = os.path.join(target_dir, DIR_LABELS_DS)
|
||||
else:
|
||||
new_file_path = os.path.join(target_dir, DIR_RGB_DS)
|
||||
|
||||
line = os.path.join("./", DIR_RGB_DS, new_file_name) + "\n"
|
||||
if nn % SZ_SERIES == 0:
|
||||
f2.write(line)
|
||||
else:
|
||||
f1.write(line)
|
||||
|
||||
# Полные пути к старому и новому файлам
|
||||
old_file_path = os.path.join(source_dir, file)
|
||||
new_file_path = os.path.join(new_file_path, new_file_name)
|
||||
|
||||
# Копируем файл
|
||||
shutil.copy2(old_file_path, new_file_path)
|
||||
|
||||
return file_nn
|
||||
|
||||
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:
|
||||
|
|
@ -67,12 +118,12 @@ def gt_parse(path: str, out_dir: str):
|
|||
# формат: <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)
|
||||
nn_image += 1
|
||||
|
||||
def explore(path: str, res_dir: str):
|
||||
if not os.path.isdir(path):
|
||||
|
|
@ -88,7 +139,7 @@ def explore(path: str, res_dir: str):
|
|||
else:
|
||||
explore(path_entry, res_dir)
|
||||
|
||||
def BOP2Yolo_dataset(dpath: str, out_dir: str, lname: list) -> str:
|
||||
def BOP2Yolo_dataset(dpath: str, out_dir: str, lname: list, addon:str) -> 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)
|
||||
|
|
@ -116,12 +167,14 @@ def BOP2Yolo_dataset(dpath: str, out_dir: str, lname: list) -> str:
|
|||
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)
|
||||
if addon:
|
||||
add_on_dataset(addon, res_dir)
|
||||
f1.close()
|
||||
f2.close()
|
||||
|
||||
return out_dir
|
||||
|
||||
def train_YoloV8(path:str, wname:str, dname:str, outpath:str, epochs:int, pretrain: bool):
|
||||
def train_YoloV8(path:str, wname:str, dname:str, outpath:str, epochs:int, pretrain: bool, addon: str):
|
||||
""" Main procedure for train YOLOv8 model """
|
||||
if not os.path.isdir(outpath):
|
||||
print(f"Invalid output path '{outpath}'")
|
||||
|
|
@ -151,21 +204,33 @@ def train_YoloV8(path:str, wname:str, dname:str, outpath:str, epochs:int, pretra
|
|||
# список имён объектов
|
||||
list_name = list(map(lambda x: x["name"], y))
|
||||
|
||||
dpath = BOP2Yolo_dataset(ds_path, out_dir, list_name)
|
||||
dpath = BOP2Yolo_dataset(ds_path, out_dir, list_name, addon)
|
||||
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)
|
||||
|
||||
# # Update settings
|
||||
# settings.update({"profile": True})
|
||||
|
||||
# prof = torch.profiler.profile(
|
||||
# schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
|
||||
# on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
|
||||
# record_shapes=True,
|
||||
# with_stack=True)
|
||||
# prof.start()
|
||||
results = model.train(data=os.path.join(dpath, FILE_RBS_TRAIN), epochs=epochs, project=out_dir) #, log_dir="runs/train")
|
||||
# prof.stop()
|
||||
|
||||
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)
|
||||
# shutil.rmtree(results.save_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
|
@ -176,6 +241,7 @@ if __name__ == "__main__":
|
|||
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")
|
||||
parser.add_argument("--addon", default="", help="Folder with add-on for dataset")
|
||||
args = parser.parse_args()
|
||||
|
||||
train_YoloV8(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain)
|
||||
train_YoloV8(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain, args.addon)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue