user3142067
user3142067

Reputation: 1292

Numpy: creating batch of numpy arrays within another numpy array (reshaping)

I have a numpy array batch of shape (32,5). Each element of the batch consists of a numpy array batch_elem = [s,_,_,_,_] where s = [img,val1,val2] is a 3-dimensional numpy array and _ are simply scalar values. img is an image (numpy array) with dimensions (84,84,3)

I would like to create a numpy array with the shape (32,84,84,3). Basically I want to extract the image information within each batch and transform it into a 4-dimensional array.

I tried the following:

b = np.vstack(batch[:,0]) #this yields a b with shape (32,3), type: <class 'numpy.ndarray'>

Now I would like to access the images (first index in second dimension)

img_batch = b[:,0] # this returns an array of shape (32,), type: <class 'numpy.ndarray'>

How can I best access the image data and get a shape (32,84,84,3)?

Note:

 s = b[0] #first s of the 32 in batch: shape (3,) , type: <class 'numpy.ndarray'>

Edit:

This should be a minimal example:

img = np.zeros([5,5,3])
s = np.array([img,1,1])
batch_elem = np.array([s,1,1,1,1])
batch = np.array([batch_elem for _ in range(32)])

Upvotes: 4

Views: 5265

Answers (2)

hilberts_drinking_problem
hilberts_drinking_problem

Reputation: 11602

Assuming I understand the problem correctly, you can stack twice on the last axis.

res = np.stack(np.stack(batch[:,0])[...,0])

Upvotes: 3

Sanghack Lee
Sanghack Lee

Reputation: 101

import numpy as np

# fabricate some data
batch = np.array((32, 1), dtype=object)
for i in range(len(batch)):
    batch[i] = [np.random.rand(84, 84, 3), None, None]

# select images
result = np.array([img for img, _, _ in batch])

# double check!
for i in range(len(batch)):
    assert np.all(result[i, :, :, :] == batch[i][0])

Upvotes: 1

Related Questions