How to loop through a nested list, compare the first element to another list, and then append to a new list?

Viewed 65

I am trying to create three separate lists (training, testing, and validation) for a CNN.
I have three lists (train_data, test_data, val_data) that ONLY contain the image names

train_data = ["img_1.png", "img_2.png"] 
test_data = ["image_3.png", "img_4.png"] 
val_data = ["img_5.png", "img_6.png"]

I have another nested list that contains all the images names and associated labels for a deep learning model

image_annotations = [['img_1.png', 432, 662, 554, 749, 'class'], 
    ['img_1.png', 647, 456, 754, 594, 'class'], ['img_2.png', 598, 659, 897, 
    302, 'class']]

I would like to run a loop through my image_annotations list and if the image name is the same as an image name in my train_data, test_data, val_data lists then I want to append the four following numbers and class name to that list. The image_annotations list has multiple listings of the same image but with different bounding box numbers. I imagine that this might be quite simple, but I can't figure it out.

3 Answers

I think this may help you:

train_data_full = []
for t_d in train_data:
    for i_a in image_annotations:
        if t_d in i_a:
            train_data_full.append(i_a)

This would be the answer to your question

train_data = ["img_1.png", "img_2.png"] 
test_data = ["image_3.png", "img_4.png"] 
val_data = ["img_5.png", "img_6.png"]
image_annotations = [['img_1.png', 432, 662, 554, 749, 'class'],['img_1.png', 647, 456, 754, 594, 'class'],['img_2.png', 598, 659, 897, 302, 'class']]
my_lists = [train_data, test_data, val_data]
for im in image_annotations:
    for k in my_lists:
        if im[0] in k:
            for i in range(1,5):
                k.append(im[i])

for list in my_lists:
    print(list)

you can use:

train_data = ["img_1.png", "img_2.png"] 
test_data = ["image_3.png", "img_4.png"] 
val_data = ["img_5.png", "img_6.png"]

image_annotations = [['img_1.png', 432, 662, 554, 749, 'class'], 
    ['img_1.png', 647, 456, 754, 594, 'class'], ['img_2.png', 598, 659, 897, 
    302, 'class']]

# get a maping with all the img names and their value
d = {}
for e in image_annotations:
    d.setdefault(e[0], []).append(e)


# set new values to data variables accordding to dict d
for l in train_data, test_data, val_data:
    l[:] = [e for i in l for e in d.get(i, [i])]


print(train_data)
print(test_data)
print(val_data)

output:

[['img_1.png', 432, 662, 554, 749, 'class'], ['img_1.png', 647, 456, 754, 594, 'class'], ['img_2.png', 598, 659, 897, 302, 'class']]
['image_3.png', 'img_4.png']
['img_5.png', 'img_6.png']
Related