import mxnet as mx
import mxnet.ndarray as nd
from skimage import io
import numpy as np
rec_path = path_prefix + ".rec"
idx_path = path_prefix + ".idx"
print(rec_path)

train_iter = mx.image.ImageIter(
    batch_size=3,
    data_shape=(3, 512, 512),
    path_imgrec=rec_path,
    path_imgidx=idx_path,
    shuffle=True)  

train_iter.reset()

for batch in train_iter:
    x = batch.data[0]
    y = batch.label[0]
    break

print(y[0])

img = nd.transpose(x, (0, 2, 3, 1))
print(img.shape)
io.imshow(img[0].asnumpy().astype(np.uint8))
io.show()