Returning categories in sample with DataLoader for Coco dataset

1

I am looking to do semantic segmentation using the coco dataset. The provided DataLoader does a good job giving me my images and annotations, but I would also like to return a y_cls 1D tensor, where y_cls[i, T] = 1 if class T is present in an image for each iteration.

I was thinking of generating y_cls something like:


cats = [1, 2, 3, 4, 5]
y_ann_labels = [{'id': 1}, {'id': 3}, {'id': 3}, {'id': 4}]
y_cls = []
temp = []

for i in y_ann_labels:
    temp.append(i.get('id'))

for i in cats:
        if i in temp:
            y_cls.append(1)
        else: 
            y_cls.append(0)
print y_cls
[1, 0, 1, 1, 0]

# then make y_cls into a tensor

I’m not sure how to combine this with the DataLoader though, so that I am returning x, y, y_cls for every iteration. Any suggestions?