Get random element on condition

Question

I have two lists: images and corresponding labels. I want to get a random image by a given label.

I thought about using numpy array to get a boolean array and then use list.index. The problem is that it returns the index of the first occurance. Any ideas?

Answer 1

Use numpy.where to get an array containing the indices where the condition is true and then use numpy.random.choice to randomly choose one:

import numpy as np
# to always get the same result
np.random.seed(42)

# setup data
images = np.arange(10)
labels = np.arange(10)

# select on the labels
indices = np.where(labels % 2 == 0)[0]
print(indices)
# [0, 2, 4, 6, 8]

# choose one
random_image = images[np.random.choice(indices)]
print(random_image)
# 6

You probably want to put that in a function:

from numpy import where
from numpy.random import choice

def random_img_with_label(images, labels, label):
    indices = where(labels == label)[0]
    return images[choice(indices)]

Alternatively you can directly create a mask from your condition and choose from that:

random_image = np.random.choice(images[labels % 2 == 0])