Validate features during add_frame + Add 2D-to-5D + Add string

[Cadene]

What this does

Motivations: We want to support more modalities and data type that Hugging Face datasets already supports. See: https://github.com/huggingface/datasets/blob/339e9dc3914e4ed837f2c9b24972e211334025ec/src/datasets/features/features.py#L462-L646

• Add ndarray features of shape: (1,), (n,), (n,m), (n,m,l), (n,m,l,o), (n,m,l,o,p)
• Add string features like caption per frame
• Add sanity check at the top of add_frame to validate if frame matches self.features specification
  • Display informative error messages when some features are missing, or extra features are present, the dtype is wrong, the shape is wrong.
• Add range check to image_array_to_pil_image which is executed in the image writer threads

How it was tested

• Added a lot of tests for add_frame in pytest -sx tests/test_datasets.py
• Added tests for image_array_to_pil_image with range_check=True.
• The range check involves the compute of a max and a min which are fast enough: Average time taken for computing max and min over 100 runs: 0.000188 seconds

import numpy as np
import time

def main():
    array = np.random.rand(3,480,640)
    num_repetitions = 100
    elapsed_times = []
    for _ in range(num_repetitions):
        # Measure the time taken to perform torch.max
        start_time = time.perf_counter()
        array.max()
        array.min()
        end_time = time.perf_counter()

        # Calculate the elapsed time
        elapsed_time = end_time - start_time
        elapsed_times.append(elapsed_time)

    # Calculate the average elapsed time
    average_elapsed_time = sum(elapsed_times) / num_repetitions
    print(f"Average time taken for computing max and min over {num_repetitions} runs: {average_elapsed_time:.6f} seconds")

if __name__ == "__main__":
    main()

[aliberts]

First round, let's iterate on the comments a bit ;)

[aliberts]

LGTM

[aliberts]

Note: This causes any string features (such as "caption") to be stored as string inside np.arrays which is really not ideal. It cause failing tests after rebasing these changes on the per-episodes stats branch (see hack). We should probably just continue in these cases, WDYT @Cadene?