DSPmagician
DSPmagician

Reputation: 81

How to select only one part of a Tensorflow dataset, and change the dimensions

I wish to train my model on 10 frame segments of UCF101, without any label. Currently I have this:

import tensorflow as tf
import tensorflow_datasets as tfds
x_train = tfds.load('ucf101', split='train', shuffle_files=True, batch_size = 64)
>>> print(x_train)
<_OptionsDataset shapes: {label: (None,), video: (None, None, 256, 256, 3)}, types: {label: tf.int64, video: tf.uint8}>

I would like the dimensions of the dataset to be (None, 10, 256, 256, 3), and not include the label.

Edit: I tried using lambda expressions in .map(), but this yielded an error.

new_x_train = x_train.map(lambda x: tf.py_function(func=lambda y: tf.convert_to_tensor(sample(y.numpy().tolist(), 10), dtype=uint8), inp=[x['video']], Tout=tf.uint8))
NameError: name 'sample' is not defined

Upvotes: 2

Views: 2083

Answers (2)

DSPmagician
DSPmagician

Reputation: 81

The solution to this was simply to download the dataset files elsewhere, so I had a list of .avi files in my directory, and then preprocess these files outside of tensorflow. I used the cv2 library and the following code, where I borrowed the two functions from elsewhere:

# Utilities to open video files using CV2
def crop_center_square(frame):
  y, x = frame.shape[0:2]
  min_dim = min(y, x)
  start_x = (x // 2) - (min_dim // 2)
  start_y = (y // 2) - (min_dim // 2)
  return frame[start_y:start_y+min_dim,start_x:start_x+min_dim]

def load_video(path, max_frames=0, resize=(256, 256)):
  cap = cv2.VideoCapture(path)
  frames = []
  try:
    while True:
      ret, frame = cap.read()
      if not ret:
        break
      frame = crop_center_square(frame)
      frame = cv2.resize(frame, resize)
      frame = frame[:, :, [2, 1, 0]]
      frames.append(frame)

      if len(frames) == max_frames:
        break
  finally:
    cap.release()
  return np.array(frames) / 255.0


files = [f for f in glob.glob("**/*.avi", recursive=True)]

for video_path in files:
  video = load_video(video_path)
  video_name = video_path[video_path.find('/')+1:]
  num_frames = video.shape[0]
  print("Video in " + video_path + " has " + str(num_frames) + " frames.")
  for seg_num in range(math.floor(num_frames/10)):
    result = video[seg_num*10:(seg_num+1)*10, ...]
    new_filepath = video_name[:-4] + "_" + str(seg_num).zfill(2) + ".avi"
    print(new_filepath)
    out = cv2.VideoWriter(new_filepath,0, 25.0, (256,256))
    for frame_n in range(0,10):
      out.write(np.uint8(255*result[frame_n, ...]))
    out.release()
    del result
  del video

Upvotes: 1

Nicolas Gervais
Nicolas Gervais

Reputation: 36704

Forgive me for the approximate answer, because I won't download the 6GB dataset to test my answer.

Why don't you just select the video when you iterate through the dataset:

next(iter(x_train))['video']

To select the dimensions, you can use normal numpy indexing. That would be an example with mnist:

import tensorflow_datasets as tfds

data = tfds.load('mnist', split='train', batch_size=16)
<PrefetchDataset shapes: {image: (None, 28, 28, 1), 
    label: (None,)}, types: {image: tf.uint8, label: tf.int64}>

Now let's select only image, and select the first 10 observations.

dim = lambda x: x['image'][:10, ...]

next(iter(data.map(dim))).shape
TensorShape([10, 28, 28, 1])

See how I removed a None in the shape with simple indexing.

Upvotes: 2

Related Questions