dataset_generation and train_models as modules

train_models/train_Dope.py

@@ -0,0 +1,542 @@
+"""
+    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)