543 lines
20 KiB
Python
543 lines
20 KiB
Python
|
"""
|
|||
|
|
train_Dope
|
|||
|
|
Общая задача: оценка позиции объекта (Pose estimation)
|
|||
|
|
Реализуемая функция: обучение нейросетевой модели DOPE по заданному BOP-датасету
|
|||
|
|
|
|||
|
|
python3 $PYTHON_EDUCATION --path /Users/user/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \
|
|||
|
|
--name test1234 --datasetName 32123213
|
|||
|
|
|
|||
|
|
08.05.2024 @shalenikol release 0.1
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
import shutil
|
|||
|
|
import numpy as np
|
|||
|
|
import transforms3d as t3d
|
|||
|
|
|
|||
|
|
FILE_RBS_INFO = "rbs_info.json"
|
|||
|
|
FILE_CAMERA = "camera.json"
|
|||
|
|
FILE_GT = "scene_gt.json"
|
|||
|
|
FILE_GT_COCO = "scene_gt_coco.json"
|
|||
|
|
FILE_GT_INFO = "scene_gt_info.json"
|
|||
|
|
|
|||
|
|
FILE_MODEL = "epoch"
|
|||
|
|
EXT_MODEL = ".pth"
|
|||
|
|
EXT_RGB = "jpg"
|
|||
|
|
DIR_ROOT_DS = "dataset_dope"
|
|||
|
|
DIR_TRAIN_OUT = "out_weights"
|
|||
|
|
|
|||
|
|
MODEL_SCALE = 1000 # исходная модель в метрах, преобразуем в мм (для DOPE)
|
|||
|
|
|
|||
|
|
# Own_Numbering_Files = True # наименование image-файлов: собственная нумерация
|
|||
|
|
nn_image = 0
|
|||
|
|
K_intrinsic = []
|
|||
|
|
model_info = []
|
|||
|
|
camera_data = {}
|
|||
|
|
im_width = 0
|
|||
|
|
|
|||
|
|
nb_update_network = 0
|
|||
|
|
# [
|
|||
|
|
# [min(x), min(y), min(z)],
|
|||
|
|
# [min(x), max(y), min(z)],
|
|||
|
|
# [min(x), max(y), max(z)],
|
|||
|
|
# [min(x), min(y), max(z)],
|
|||
|
|
# [max(x), min(y), max(z)],
|
|||
|
|
# [max(x), max(y), min(z)],
|
|||
|
|
# [max(x), max(y), max(z)],
|
|||
|
|
# [max(x), min(y), max(z)],
|
|||
|
|
# [xc, yc, zc] # min + (max - min) / 2
|
|||
|
|
# ]
|
|||
|
|
|
|||
|
|
def trans_3Dto2D_point_in_camera(xyz, K_m, R_m2c, t_m2c):
|
|||
|
|
"""
|
|||
|
|
xyz : 3D-координаты точки
|
|||
|
|
K_m : внутренняя матрица камеры 3х3
|
|||
|
|
R_m2c : матрица поворота 3х3
|
|||
|
|
t_m2c : вектор перемещения 3х1
|
|||
|
|
return [u,v]
|
|||
|
|
"""
|
|||
|
|
K = np.array(K_m)
|
|||
|
|
r = np.array(R_m2c)
|
|||
|
|
r.shape = (3, 3)
|
|||
|
|
t = np.array(t_m2c)
|
|||
|
|
t.shape = (3, 1)
|
|||
|
|
T = np.concatenate((r, t), axis=1)
|
|||
|
|
|
|||
|
|
P_m = np.array(xyz)
|
|||
|
|
P_m.resize(4)
|
|||
|
|
P_m[-1] = 1.0
|
|||
|
|
P_m.shape = (4, 1)
|
|||
|
|
|
|||
|
|
# Project (X, Y, Z, 1) into cameras coordinate system
|
|||
|
|
P_c = T @ P_m # 4x1
|
|||
|
|
# Apply camera intrinsics to map (Xc, Yc, Zc) to p=(x, y, z)
|
|||
|
|
p = K @ P_c
|
|||
|
|
# Normalize by z to get (u,v,1)
|
|||
|
|
uv = (p / p[2][0])[:-1]
|
|||
|
|
return uv.flatten().tolist()
|
|||
|
|
|
|||
|
|
def gt_parse(path: str, out_dir: str):
|
|||
|
|
global nn_image
|
|||
|
|
with open(os.path.join(path, FILE_GT_COCO), "r") as fh:
|
|||
|
|
coco_data = json.load(fh)
|
|||
|
|
with open(os.path.join(path, FILE_GT), "r") as fh:
|
|||
|
|
gt_data = json.load(fh)
|
|||
|
|
with open(os.path.join(path, FILE_GT_INFO), "r") as fh:
|
|||
|
|
gt_info = json.load(fh)
|
|||
|
|
|
|||
|
|
for img in coco_data["images"]:
|
|||
|
|
rgb_file = os.path.join(path, img["file_name"])
|
|||
|
|
if os.path.isfile(rgb_file):
|
|||
|
|
# if Own_Numbering_Files:
|
|||
|
|
ext = os.path.splitext(rgb_file)[1] # only ext
|
|||
|
|
f = f"{nn_image:06}"
|
|||
|
|
out_img = os.path.join(out_dir, f + ext)
|
|||
|
|
# else:
|
|||
|
|
# f = os.path.split(rgb_file)[1] # filename with extension
|
|||
|
|
# f = os.path.splitext(f)[0] # only filename
|
|||
|
|
# out_img = out_dir
|
|||
|
|
shutil.copy2(rgb_file, out_img)
|
|||
|
|
out_file = os.path.join(out_dir,f+".json")
|
|||
|
|
nn_image += 1
|
|||
|
|
|
|||
|
|
# full annotation of the one image
|
|||
|
|
all_data = camera_data.copy()
|
|||
|
|
cat_names = {obj["id"]: obj["name"] for obj in coco_data["categories"]}
|
|||
|
|
id_img = img["id"] # 0, 1, 2 ...
|
|||
|
|
sid_img = str(id_img) # "0", "1", "2" ...
|
|||
|
|
img_info = gt_info[sid_img]
|
|||
|
|
img_gt = gt_data[sid_img]
|
|||
|
|
img_idx = 0 # object index on the image
|
|||
|
|
objs = []
|
|||
|
|
for ann in coco_data["annotations"]:
|
|||
|
|
if ann["image_id"] == id_img:
|
|||
|
|
item = ann["category_id"]
|
|||
|
|
obj_data = {}
|
|||
|
|
obj_data["class"] = cat_names[item]
|
|||
|
|
x, y, width, height = ann["bbox"]
|
|||
|
|
obj_data["bounding_box"] = {"top_left":[x,y], "bottom_right":[x+width,y+height]}
|
|||
|
|
|
|||
|
|
# visibility from FILE_GT_INFO
|
|||
|
|
item_info = img_info[img_idx]
|
|||
|
|
obj_data["visibility"] = item_info["visib_fract"]
|
|||
|
|
|
|||
|
|
# location from FILE_GT
|
|||
|
|
item_gt = img_gt[img_idx]
|
|||
|
|
obj_id = item_gt["obj_id"] - 1 # index with 0
|
|||
|
|
cam_R_m2c = item_gt["cam_R_m2c"]
|
|||
|
|
cam_t_m2c = item_gt["cam_t_m2c"]
|
|||
|
|
obj_data["location"] = cam_t_m2c
|
|||
|
|
q = t3d.quaternions.mat2quat(np.array(cam_R_m2c))
|
|||
|
|
obj_data["quaternion_xyzw"] = [q[1], q[2], q[3], q[0]]
|
|||
|
|
|
|||
|
|
cuboid_xyz = model_info[obj_id]
|
|||
|
|
obj_data["projected_cuboid"] = [
|
|||
|
|
trans_3Dto2D_point_in_camera(cub, K_intrinsic, cam_R_m2c, cam_t_m2c)
|
|||
|
|
for cub in cuboid_xyz
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
objs.append(obj_data)
|
|||
|
|
img_idx += 1
|
|||
|
|
|
|||
|
|
all_data["objects"] = objs
|
|||
|
|
with open(out_file, "w") as fh:
|
|||
|
|
json.dump(all_data, fh, indent=2)
|
|||
|
|
|
|||
|
|
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)) and \
|
|||
|
|
os.path.isfile(os.path.join(path_entry,FILE_GT_INFO)) and \
|
|||
|
|
os.path.isfile(os.path.join(path_entry,FILE_GT)):
|
|||
|
|
gt_parse(path_entry, res_dir)
|
|||
|
|
else:
|
|||
|
|
explore(path_entry, res_dir)
|
|||
|
|
|
|||
|
|
def BOP2DOPE_dataset(dpath: str, out_dir: str) -> str:
|
|||
|
|
""" Convert BOP-dataset to YOLO format for train """
|
|||
|
|
res_dir = os.path.join(out_dir, DIR_ROOT_DS)
|
|||
|
|
if os.path.isdir(res_dir):
|
|||
|
|
shutil.rmtree(res_dir)
|
|||
|
|
os.mkdir(res_dir)
|
|||
|
|
|
|||
|
|
explore(dpath, res_dir)
|
|||
|
|
|
|||
|
|
return out_dir
|
|||
|
|
|
|||
|
|
def train(dopepath:str, wname:str, epochs:int, pretrain: bool, lname: list):
|
|||
|
|
import random
|
|||
|
|
# try:
|
|||
|
|
import configparser as configparser
|
|||
|
|
# except ImportError:
|
|||
|
|
# import ConfigParser as configparser
|
|||
|
|
import torch
|
|||
|
|
# import torch.nn.parallel
|
|||
|
|
import torch.optim as optim
|
|||
|
|
import torch.utils.data
|
|||
|
|
import torchvision.transforms as transforms
|
|||
|
|
from torch.autograd import Variable
|
|||
|
|
import datetime
|
|||
|
|
from tensorboardX import SummaryWriter
|
|||
|
|
|
|||
|
|
from models_dope import DopeNetwork
|
|||
|
|
from utils_dope import CleanVisiiDopeLoader #, VisualizeBeliefMap, save_image
|
|||
|
|
|
|||
|
|
import warnings
|
|||
|
|
warnings.filterwarnings("ignore")
|
|||
|
|
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
|
|||
|
|
|
|||
|
|
torch.autograd.set_detect_anomaly(False)
|
|||
|
|
torch.autograd.profiler.profile(False)
|
|||
|
|
torch.autograd.gradcheck = False
|
|||
|
|
torch.backends.cudnn.benchmark = True
|
|||
|
|
|
|||
|
|
start_time = datetime.datetime.now()
|
|||
|
|
print("start:", start_time.strftime("%m/%d/%Y, %H:%M:%S"))
|
|||
|
|
|
|||
|
|
res_model = os.path.join(dopepath, wname + EXT_MODEL)
|
|||
|
|
|
|||
|
|
local_rank = 0
|
|||
|
|
opt = lambda: None
|
|||
|
|
opt.use_s3 = False
|
|||
|
|
opt.train_buckets = []
|
|||
|
|
opt.endpoint = None
|
|||
|
|
opt.lr=0.0001
|
|||
|
|
opt.loginterval=100
|
|||
|
|
opt.sigma=0.5 # 4
|
|||
|
|
opt.nbupdates=None
|
|||
|
|
# opt.save=False
|
|||
|
|
# opt.option="default"
|
|||
|
|
# opt.gpuids=[0]
|
|||
|
|
|
|||
|
|
opt.namefile=FILE_MODEL
|
|||
|
|
opt.workers=8
|
|||
|
|
opt.batchsize=16
|
|||
|
|
|
|||
|
|
opt.data = [os.path.join(dopepath, DIR_ROOT_DS)]
|
|||
|
|
opt.outf = os.path.join(dopepath, DIR_TRAIN_OUT)
|
|||
|
|
opt.object = lname #["fork"]
|
|||
|
|
opt.exts = [EXT_RGB]
|
|||
|
|
# opt.imagesize = im_width
|
|||
|
|
opt.epochs = epochs
|
|||
|
|
opt.pretrained = pretrain
|
|||
|
|
opt.net_path = res_model if pretrain else None
|
|||
|
|
opt.manualseed = random.randint(1, 10000)
|
|||
|
|
|
|||
|
|
# # Validate Arguments
|
|||
|
|
# if opt.use_s3 and (opt.train_buckets is None or opt.endpoint is None):
|
|||
|
|
# raise ValueError(
|
|||
|
|
# "--train_buckets and --endpoint must be specified if training with data from s3 bucket."
|
|||
|
|
# )
|
|||
|
|
# if not opt.use_s3 and opt.data is None:
|
|||
|
|
# raise ValueError("--data field must be specified.")
|
|||
|
|
|
|||
|
|
os.makedirs(opt.outf, exist_ok=True)
|
|||
|
|
|
|||
|
|
# if local_rank == 0:
|
|||
|
|
# writer = SummaryWriter(opt.outf + "/runs/")
|
|||
|
|
random.seed(opt.manualseed)
|
|||
|
|
torch.cuda.set_device(local_rank)
|
|||
|
|
# torch.distributed.init_process_group(backend="nccl", init_method="env://")
|
|||
|
|
torch.manual_seed(opt.manualseed)
|
|||
|
|
torch.cuda.manual_seed_all(opt.manualseed)
|
|||
|
|
|
|||
|
|
# # Data Augmentation
|
|||
|
|
# if not opt.save:
|
|||
|
|
# contrast = 0.2
|
|||
|
|
# brightness = 0.2
|
|||
|
|
# noise = 0.1
|
|||
|
|
# normal_imgs = [0.59, 0.25]
|
|||
|
|
# transform = transforms.Compose(
|
|||
|
|
# [
|
|||
|
|
# AddRandomContrast(0.2),
|
|||
|
|
# AddRandomBrightness(0.2),
|
|||
|
|
# transforms.Resize(opt.imagesize),
|
|||
|
|
# ]
|
|||
|
|
# )
|
|||
|
|
# else:
|
|||
|
|
# contrast = 0.00001
|
|||
|
|
# brightness = 0.00001
|
|||
|
|
# noise = 0.00001
|
|||
|
|
# normal_imgs = None
|
|||
|
|
# transform = transforms.Compose(
|
|||
|
|
# [transforms.Resize(opt.imagesize), transforms.ToTensor()]
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
# Load Model
|
|||
|
|
net = DopeNetwork()
|
|||
|
|
output_size = 50
|
|||
|
|
# opt.sigma = 0.5
|
|||
|
|
|
|||
|
|
train_dataset = CleanVisiiDopeLoader(
|
|||
|
|
opt.data,
|
|||
|
|
sigma=opt.sigma,
|
|||
|
|
output_size=output_size,
|
|||
|
|
extensions=opt.exts,
|
|||
|
|
objects=opt.object,
|
|||
|
|
use_s3=opt.use_s3,
|
|||
|
|
buckets=opt.train_buckets,
|
|||
|
|
endpoint_url=opt.endpoint,
|
|||
|
|
)
|
|||
|
|
trainingdata = torch.utils.data.DataLoader(
|
|||
|
|
train_dataset,
|
|||
|
|
batch_size=opt.batchsize,
|
|||
|
|
shuffle=True,
|
|||
|
|
num_workers=opt.workers,
|
|||
|
|
pin_memory=True,
|
|||
|
|
)
|
|||
|
|
if not trainingdata is None:
|
|||
|
|
print(f"training data: {len(trainingdata)} batches")
|
|||
|
|
|
|||
|
|
print("Loading Model...")
|
|||
|
|
net = net.cuda()
|
|||
|
|
# net = torch.nn.parallel.DistributedDataParallel(
|
|||
|
|
# net.cuda(), device_ids=[local_rank], output_device=local_rank
|
|||
|
|
# )
|
|||
|
|
if opt.pretrained:
|
|||
|
|
if opt.net_path is not None:
|
|||
|
|
net.load_state_dict(torch.load(opt.net_path))
|
|||
|
|
else:
|
|||
|
|
print("Error: Did not specify path to pretrained weights.")
|
|||
|
|
quit()
|
|||
|
|
|
|||
|
|
parameters = filter(lambda p: p.requires_grad, net.parameters())
|
|||
|
|
optimizer = optim.Adam(parameters, lr=opt.lr)
|
|||
|
|
|
|||
|
|
print("ready to train!")
|
|||
|
|
|
|||
|
|
global nb_update_network
|
|||
|
|
nb_update_network = 0
|
|||
|
|
# best_results = {"epoch": None, "passed": None, "add_mean": None, "add_std": None}
|
|||
|
|
|
|||
|
|
scaler = torch.cuda.amp.GradScaler()
|
|||
|
|
|
|||
|
|
def _runnetwork(epoch, train_loader): #, syn=False
|
|||
|
|
global nb_update_network
|
|||
|
|
# net
|
|||
|
|
net.train()
|
|||
|
|
|
|||
|
|
loss_avg_to_log = {}
|
|||
|
|
loss_avg_to_log["loss"] = []
|
|||
|
|
loss_avg_to_log["loss_affinities"] = []
|
|||
|
|
loss_avg_to_log["loss_belief"] = []
|
|||
|
|
loss_avg_to_log["loss_class"] = []
|
|||
|
|
for batch_idx, targets in enumerate(train_loader):
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
|
|||
|
|
data = Variable(targets["img"].cuda())
|
|||
|
|
target_belief = Variable(targets["beliefs"].cuda())
|
|||
|
|
target_affinities = Variable(targets["affinities"].cuda())
|
|||
|
|
|
|||
|
|
output_belief, output_aff = net(data)
|
|||
|
|
|
|||
|
|
loss = None
|
|||
|
|
|
|||
|
|
loss_belief = torch.tensor(0).float().cuda()
|
|||
|
|
loss_affinities = torch.tensor(0).float().cuda()
|
|||
|
|
loss_class = torch.tensor(0).float().cuda()
|
|||
|
|
|
|||
|
|
for stage in range(len(output_aff)): # output, each belief map layers.
|
|||
|
|
loss_affinities += (
|
|||
|
|
(output_aff[stage] - target_affinities)
|
|||
|
|
* (output_aff[stage] - target_affinities)
|
|||
|
|
).mean()
|
|||
|
|
|
|||
|
|
loss_belief += (
|
|||
|
|
(output_belief[stage] - target_belief)
|
|||
|
|
* (output_belief[stage] - target_belief)
|
|||
|
|
).mean()
|
|||
|
|
|
|||
|
|
loss = loss_affinities + loss_belief
|
|||
|
|
|
|||
|
|
# if batch_idx == 0:
|
|||
|
|
# post = "train"
|
|||
|
|
# if local_rank == 0:
|
|||
|
|
# for i_output in range(1):
|
|||
|
|
# # input images
|
|||
|
|
# writer.add_image(
|
|||
|
|
# f"{post}_input_{i_output}",
|
|||
|
|
# targets["img_original"][i_output],
|
|||
|
|
# epoch,
|
|||
|
|
# dataformats="CWH",
|
|||
|
|
# )
|
|||
|
|
# # belief maps gt
|
|||
|
|
# imgs = VisualizeBeliefMap(target_belief[i_output])
|
|||
|
|
# img, grid = save_image(
|
|||
|
|
# imgs, "some_img.png", mean=0, std=1, nrow=3, save=False
|
|||
|
|
# )
|
|||
|
|
# writer.add_image(
|
|||
|
|
# f"{post}_belief_ground_truth_{i_output}",
|
|||
|
|
# grid,
|
|||
|
|
# epoch,
|
|||
|
|
# dataformats="CWH",
|
|||
|
|
# )
|
|||
|
|
# # belief maps guess
|
|||
|
|
# imgs = VisualizeBeliefMap(output_belief[-1][i_output])
|
|||
|
|
# img, grid = save_image(
|
|||
|
|
# imgs, "some_img.png", mean=0, std=1, nrow=3, save=False
|
|||
|
|
# )
|
|||
|
|
# writer.add_image(
|
|||
|
|
# f"{post}_belief_guess_{i_output}",
|
|||
|
|
# grid,
|
|||
|
|
# epoch,
|
|||
|
|
# dataformats="CWH",
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
loss.backward()
|
|||
|
|
|
|||
|
|
optimizer.step()
|
|||
|
|
|
|||
|
|
nb_update_network += 1
|
|||
|
|
|
|||
|
|
# log the loss
|
|||
|
|
loss_avg_to_log["loss"].append(loss.item())
|
|||
|
|
loss_avg_to_log["loss_class"].append(loss_class.item())
|
|||
|
|
loss_avg_to_log["loss_affinities"].append(loss_affinities.item())
|
|||
|
|
loss_avg_to_log["loss_belief"].append(loss_belief.item())
|
|||
|
|
|
|||
|
|
if batch_idx % opt.loginterval == 0:
|
|||
|
|
print(
|
|||
|
|
"Train Epoch: {} [{}/{} ({:.0f}%)] \tLoss: {:.15f} \tLocal Rank: {}".format(
|
|||
|
|
epoch,
|
|||
|
|
batch_idx * len(data),
|
|||
|
|
len(train_loader.dataset),
|
|||
|
|
100.0 * batch_idx / len(train_loader),
|
|||
|
|
loss.item(),
|
|||
|
|
local_rank,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
# # log the loss values
|
|||
|
|
# if local_rank == 0:
|
|||
|
|
# writer.add_scalar("loss/train_loss", np.mean(loss_avg_to_log["loss"]), epoch)
|
|||
|
|
# writer.add_scalar("loss/train_cls", np.mean(loss_avg_to_log["loss_class"]), epoch)
|
|||
|
|
# writer.add_scalar("loss/train_aff", np.mean(loss_avg_to_log["loss_affinities"]), epoch)
|
|||
|
|
# writer.add_scalar("loss/train_bel", np.mean(loss_avg_to_log["loss_belief"]), epoch)
|
|||
|
|
|
|||
|
|
for epoch in range(1, opt.epochs + 1):
|
|||
|
|
|
|||
|
|
_runnetwork(epoch, trainingdata)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if local_rank == 0:
|
|||
|
|
torch.save(
|
|||
|
|
net.state_dict(),
|
|||
|
|
f"{opt.outf}/{opt.namefile}_{str(epoch).zfill(3)}.pth",
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Encountered Exception: {e}")
|
|||
|
|
|
|||
|
|
if not opt.nbupdates is None and nb_update_network > int(opt.nbupdates):
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# if local_rank == 0:
|
|||
|
|
# save result model
|
|||
|
|
torch.save(net.state_dict(), res_model) #os.path.join(dopepath, wname + EXT_MODEL))
|
|||
|
|
# else:
|
|||
|
|
# torch.save(
|
|||
|
|
# net.state_dict(),
|
|||
|
|
# f"{opt.outf}/{opt.namefile}_{str(epoch).zfill(3)}_rank_{local_rank}.pth",
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
print("end:", datetime.datetime.now().strftime("%m/%d/%Y, %H:%M:%S"))
|
|||
|
|
print("Total time taken: ", str(datetime.datetime.now() - start_time).split(".")[0])
|
|||
|
|
|
|||
|
|
def train_Dope_i(path:str, wname:str, dname:str, outpath:str, epochs:int, pretrain: bool):
|
|||
|
|
""" Main procedure for train DOPE model """
|
|||
|
|
global K_intrinsic, model_info, camera_data, im_width
|
|||
|
|
|
|||
|
|
if not os.path.isdir(outpath):
|
|||
|
|
print(f"Invalid output path '{outpath}'")
|
|||
|
|
exit(-1)
|
|||
|
|
out_dir = os.path.join(outpath, wname)
|
|||
|
|
ds_path = os.path.join(path, dname)
|
|||
|
|
|
|||
|
|
if not os.path.isdir(ds_path):
|
|||
|
|
print(f"{ds_path} : no BOP directory")
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
camera_json = os.path.join(ds_path, FILE_CAMERA)
|
|||
|
|
if not os.path.isfile(camera_json):
|
|||
|
|
print(f"{camera_json} : no intrinsic camera file")
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
rbs_info = os.path.join(ds_path, FILE_RBS_INFO)
|
|||
|
|
if not os.path.isfile(rbs_info):
|
|||
|
|
print(f"{rbs_info} : no dataset info file")
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
camera_data = {}
|
|||
|
|
with open(camera_json, "r") as fh:
|
|||
|
|
data = json.load(fh)
|
|||
|
|
keys = ["cx","cy","fx","fy"]
|
|||
|
|
intrinsic = {k: data[k] for k in keys}
|
|||
|
|
im_height = data["height"]
|
|||
|
|
im_width = data["width"]
|
|||
|
|
camera_data["camera_data"] = dict(intrinsic=intrinsic, height=im_height, width=im_width)
|
|||
|
|
K_intrinsic = [
|
|||
|
|
[data["fx"], 0.0, data["cx"]],
|
|||
|
|
[0.0, data["fy"], data["cy"]],
|
|||
|
|
[0.0, 0.0, 1.0]
|
|||
|
|
]
|
|||
|
|
# calc cuboid + center
|
|||
|
|
with open(rbs_info, "r") as fh:
|
|||
|
|
info = json.load(fh)
|
|||
|
|
# список имён объектов
|
|||
|
|
list_name = list(map(lambda x: x["name"], info))
|
|||
|
|
# in FILE_RBS_INFO model numbering from smallest to largest
|
|||
|
|
model_info = []
|
|||
|
|
for m_info in info:
|
|||
|
|
cub = np.array(m_info["cuboid"]) * MODEL_SCALE
|
|||
|
|
xyz_min = cub.min(axis=0)
|
|||
|
|
xyz_max = cub.max(axis=0)
|
|||
|
|
# [xc, yc, zc] # min + (max - min) / 2
|
|||
|
|
center = []
|
|||
|
|
for i in range(3):
|
|||
|
|
center.append(xyz_min[i] + (xyz_max[i]- xyz_min[i]) / 2)
|
|||
|
|
c = np.array(center, ndmin=2)
|
|||
|
|
model_info.append(np.append(cub, c, axis=0))
|
|||
|
|
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
dpath = BOP2DOPE_dataset(ds_path, out_dir)
|
|||
|
|
if len(dpath) == 0:
|
|||
|
|
print(f"Error in convert dataset '{ds_path}' to '{outpath}'")
|
|||
|
|
exit(-4)
|
|||
|
|
# model_path = os.path.join(dpath, FILE_BASEMODEL)
|
|||
|
|
|
|||
|
|
# results = f"python train.py --local_rank 0 --data {dpath} --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)
|
|||
|
|
train(dpath, wname, epochs, pretrain, list_name)
|
|||
|
|
|
|||
|
|
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")
|
|||
|
|
parser.add_argument('--pretrain', action="store_true", help="Use pretraining")
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain)
|