test_custom_dataloader.py

import sys, os
sys.path.insert(0, "/export/guanghan/CenterNet-Gluon/dataset")
sys.path.insert(0, "/Users/guanghan.ning/Desktop/dev/CenterNet-Gluon/dataset")
from coco_centernet import CenterCOCODataset
import mxnet as mx
from mxnet import nd, gluon, init
import gluoncv
from gluoncv.data.batchify import Tuple, Stack, Pad

def test_load():
    from opts import opts
    opt = opts().init()

    batch_size = 16
    #batchify_fn = Tuple(Stack(), Stack(), Stack(), Stack())  # stack image, heatmaps, scale, offset
    batchify_fn = Tuple(Stack(), Stack(), Stack(), Stack(), Stack(), Stack())  # stack image, heatmaps, scale, offset, ind, mask
    num_workers = 2

    train_dataset = CenterCOCODataset(opt, split = 'train')
    train_loader = gluon.data.DataLoader( train_dataset,
        batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=num_workers)
    ctx = [mx.gpu(int(i)) for i in opt.gpus_str.split(',') if i.strip()]
    ctx = ctx if ctx else [mx.cpu()]

    for i, batch in enumerate(train_loader):
        print("{} Batch".format(i))
        print("image batch shape: ", batch[0].shape)
        print("heatmap batch shape", batch[1].shape)
        print("scale batch shape", batch[2].shape)
        print("offset batch shape", batch[3].shape)
        print("indices batch shape", batch[4].shape)
        print("mask batch shape", batch[5].shape)

        X = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
        targets_heatmaps = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)  # heatmaps: (batch, num_classes, H/S, W/S)
        targets_scale = gluon.utils.split_and_load(batch[2], ctx_list=ctx, batch_axis=0)  # scale: wh (batch, 2, H/S, W/S)
        targets_offset = gluon.utils.split_and_load(batch[3], ctx_list=ctx, batch_axis=0) # offset: xy (batch, 2, H/s, W/S)
        targets_inds = gluon.utils.split_and_load(batch[4], ctx_list=ctx, batch_axis=0)
        targets_mask = gluon.utils.split_and_load(batch[5], ctx_list=ctx, batch_axis=0)

        print("len(targets_heatmaps): ", len(targets_heatmaps))
        print("First item: image shape: ", X[0].shape)
        print("First item: heatmaps shape: ", targets_heatmaps[0].shape)
        print("First item: scalemaps shape: ", targets_scale[0].shape)
        print("First item: offsetmaps shape: ", targets_offset[0].shape)
        print("First item: indices shape: ", targets_inds[0].shape)
        print("First item: mask shape: ", targets_mask[0].shape)
    return


if __name__ == "__main__":
    test_load()