Finetune on custom dataset

[Mu-Magdy]

Can you give instructions on how to prepare the data to point detection.
And how to train the model on the data?

[Kimyuhwanpeter]

@Mu-Magdy Hi, I didn't test yet but, I made custom loader based on their cocoLoader ( https://github.com/xingyizhou/CenterNet/blob/4c50fd3a46bdf63dbf2082c5cbb3458d39579e6c/src/lib/datasets/dataset/coco.py )
So right now I implemeted like this

data_config = {
"names": ["person", "cat", "dog" ],  # Class names
"nc": parser.num_classes,  # Number of classes
}
dataset = CenterBaseDataset(data=data_config, task="detect", imgsz=parser.img_size[0], 
                img_path="/home/yhkim/yuhwan_project/dataset/opencv/positive2", batch_size=parser.batch_size,
                single_cls=False, augment=True)
loader = build_dataloader(dataset, parser.batch_size, 8, True, -1)
imgsz = torch.tensor(parser.img_size, device="cuda", dtype=torch.float32)

hourglassModel.cuda()
for i, batch in pbar:
    targets = torch.cat((batch["batch_idx"].view(-1, 1), batch["cls"].view(-1, 1), batch["bboxes"]), 1)
    targets = preprocess(targets.to("cuda"), 8, scale_tensor=imgsz[[1, 0, 1, 0]], device="cuda")    # xyxy
    new_targets = GT_heatMap(targets, batch["img"])
    
    output = hourglassModel(batch["img"].float().cuda())
    hm_loss_1 = FocalLoss()(_sigmoid(output[0]["hm"]), torch.tensor(new_targets["hm"]).cuda())
    hm_loss_2 = FocalLoss()(_sigmoid(output[1]["hm"]), torch.tensor(new_targets["hm"]).cuda())
    hm_loss = (hm_loss_1 + hm_loss_2) / hourglassModel.nstack
    
    off_loss_1 = RegL1Loss()(output[0]["reg"], torch.tensor(new_targets["reg_mask"]).cuda(),
                           torch.tensor(new_targets["ind"]).cuda(), torch.tensor(new_targets["reg"]).cuda())
    off_loss_2 = RegL1Loss()(output[1]["reg"], torch.tensor(new_targets["reg_mask"]).cuda(),
                           torch.tensor(new_targets["ind"]).cuda(), torch.tensor(new_targets["reg"]).cuda())
    off_loss = (off_loss_1 + off_loss_2) / hourglassModel.nstack
    
    mask_weights = torch.tensor(new_targets["dense_wh_mask"].sum()).cuda() + 1e-4
    wh_loss_1 = torch.nn.L1Loss(reduction="sum")(output[0]["wh"] * torch.tensor(new_targets["dense_wh_mask"]).cuda(),
                                                 torch.tensor(new_targets["dense_wh"]).cuda() * torch.tensor(new_targets["dense_wh_mask"]).cuda()) / mask_weights
    wh_loss_2 = torch.nn.L1Loss(reduction="sum")(output[1]["wh"] * torch.tensor(new_targets["dense_wh_mask"]).cuda(),
                                                 torch.tensor(new_targets["dense_wh"]).cuda() * torch.tensor(new_targets["dense_wh_mask"]).cuda()) / mask_weights
    wh_loss = (wh_loss_1 + wh_loss_2) / hourglassModel.nstack
    
    total_loss = hm_loss * parser.hm_weight + off_loss * parser.off_weight + wh_loss * parser.wh_weight

and i my custom loader exclude preprocessing the heatmap, so i add it in GT_heatMap (I test the gassianMap also total_loss. in my experiment the implemented loss was fine)
maybe you can make your own loader and implemented train code like this
(this code is not include the backward)

[Kimyuhwanpeter]

and i load the " ctdet_coco_hg.pth" based hourglass model, and i inference it using my custom code
so the result was great

this is test result using ctdet_coco_hg.pth and hourglass model

inp = cv2.imread("/home/yhkim/temp.jpg")
show_img = np.array(inp, np.uint8)
OriHeight = inp.shape[0]
OriWdith = inp.shape[1]
if (OriWdith >= OriHeight):
    resize_scales = OriWdith / parser.img_size[1]
    resizedWidth = parser.img_size[1]
    resizedHeight = OriHeight / resize_scales
else:
    resize_scales = OriHeight / parser.img_size[0]
    resizedWidth = OriWdith / resize_scales
    resizedHeight = parser.img_size[0]

x_offset = (resizedWidth - parser.img_size[1]) / 2
y_offset = (resizedHeight - parser.img_size[0]) / 2


inp = testResizeImage(inp, parser.img_size[0], True)

hourglassModel.cuda()
hourglassModel.eval()
with torch.no_grad():
    output = hourglassModel(inp)

output[-1]["hm"] = _sigmoid(output[-1]["hm"])

preds = ctdet_decode(output[-1]["hm"], output[-1]["wh"], reg=output[-1]["reg"],
                    cat_spec_wh=False, K=1000)
preds[:, :, :4] *= parser.down_ratio

results = []
for i in range(preds.shape[0]):
    mask = preds[i, :, 4] > 0.25
    selected = preds[i, mask]
    results.append(selected)
preds = (results)

for pred in preds:
    pred[:, 1] = ((pred[:, 1] - math.fabs(y_offset)) / resizedHeight) * show_img.shape[0]
    pred[:, 3] = ((pred[:, 3] - math.fabs(y_offset)) / resizedHeight) * show_img.shape[0]
    pred[:, 0] = ((pred[:, 0] - math.fabs(x_offset)) / resizedWidth) * show_img.shape[1]
    pred[:, 2] = ((pred[:, 2] - math.fabs(x_offset)) / resizedWidth) * show_img.shape[1]
    pred[:, :4] = (scale_boxes((OriHeight, OriWdith), pred[:, :4], (OriHeight, OriWdith)))
    y0 = pred[:, 0: 1]  # xmin
    y1 = pred[:, 1: 2]  # ymin
    y2 = pred[:, 2: 3] - pred[:, 0: 1]  # width
    y3 = pred[:, 3: 4] - pred[:, 1: 2]  # height
    pred[:, :4] = torch.cat((y0, y1, y2, y3), dim=-1)