JPEG encoding and decoding if the observation is an image #275
Conversation
There was a problem hiding this comment.
Thanks for adding this feature!
Can you add the line .idea in the .gitignore file?
Also, can you add some tests? You can make a fake env that generates images as obs, and check that everything is okay (serialization works and deserialization returns the same obs). Check the dataset_creation tests for the examples
| @@ -157,4 +157,4 @@ cython_debug/ | |||
| # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore | |||
| # and can be added to the global gitignore or merged into this file. For a more nuclear | |||
| # option (not recommended) you can uncomment the following to ignore the entire idea folder. | |||
| #.idea/ | |||
| .idea | |||
There was a problem hiding this comment.
Uhm, maybe you should keep the '/'
The point of this is that the idea files in the code diff should disappear
| class TestEnv: | ||
| def __init__(self): | ||
| self.observation_space = spaces.Box(low=0, high=255, shape=(4, 84, 84), dtype=np.uint8) | ||
| self.action_space = spaces.Discrete(2) | ||
|
|
||
| def reset(self): | ||
| obs = np.random.randint(0, 256, (4, 84, 84), dtype=np.uint8) | ||
| return obs, {} | ||
|
|
||
| def step(self, action): | ||
| obs = np.random.randint(0, 256, (4, 84, 84), dtype=np.uint8) | ||
| reward = 1.0 | ||
| terminated = False | ||
| truncated = False | ||
| info = {} | ||
| return obs, reward, terminated, truncated, info |
There was a problem hiding this comment.
This should go in common.py in tests
And I would call it something like "ImageObsEnv"
| def generate_episode_buffer_from_env(env: TestEnv, length=3) -> EpisodeBuffer: | ||
| initial_obs , _ = env.reset() | ||
| buffer = EpisodeBuffer(observations=initial_obs) | ||
| for i in range(length): | ||
| action = env.action_space.sample() | ||
| obs, reward, terminated, truncated, info = env.step(action) | ||
| step_data = { | ||
| "observation": obs, | ||
| "action": action, | ||
| "reward": reward, | ||
| "terminated": terminated, | ||
| "truncated": truncated, | ||
| "info": info | ||
| } | ||
| buffer = buffer.add_step_data(step_data) | ||
| return buffer |
There was a problem hiding this comment.
This function is already present in common.py, just use that one
| def test_arrow_storage_serialization(): | ||
| env = TestEnv() | ||
| episode = generate_episode_buffer_from_env(env, length=3) | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| tmp_path = pathlib.Path(tmpdir) | ||
|
|
||
| METADATA_FILE_NAME = "metadata.json" | ||
| default_metadata = { | ||
| "total_steps": 0, | ||
| "total_episodes": 0, | ||
| "data_format": "arrow", | ||
| "observation_space": str(env.observation_space), | ||
| "action_space": str(env.action_space) | ||
| } | ||
| metadata_path = tmp_path.joinpath(METADATA_FILE_NAME) | ||
| metadata_path.write_text(json.dumps(default_metadata)) | ||
|
|
||
| from minari.dataset._storages.arrow_storage import ArrowStorage | ||
| storage = ArrowStorage(tmp_path, env.observation_space, env.action_space) | ||
| storage.update_episodes([episode]) | ||
| loaded_episode = list(storage.get_episodes([0])) | ||
| loaded_obs = loaded_episode[0]["observations"] | ||
| np.testing.assert_array_equal(episode.observations, loaded_obs) | ||
|
|
||
| def test_hdf5_serialization(): | ||
| env = TestEnv() | ||
| episode = generate_episode_buffer_from_env(env, length=3) | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| tmp_path = pathlib.Path(tmpdir) | ||
|
|
||
| METADATA_FILE_NAME = "metadata.json" | ||
| default_metadata = { | ||
| "total_steps": 0, | ||
| "total_episodes": 0, | ||
| "data_format": "hdf5", | ||
| "observation_space": str(env.observation_space), | ||
| "action_space": str(env.action_space) | ||
| } | ||
| metadata_path = tmp_path.joinpath(METADATA_FILE_NAME) | ||
| metadata_path.write_text(json.dumps(default_metadata)) | ||
|
|
||
| from minari.dataset._storages.hdf5_storage import HDF5Storage | ||
| storage = HDF5Storage._create(tmp_path, env.observation_space, env.action_space) | ||
| storage.update_episodes([episode]) | ||
| loaded_episodes = list(storage.get_episodes([0])) | ||
| loaded_obs = loaded_episodes[0]["observations"] | ||
| np.testing.assert_array_equal(episode.observations, loaded_obs) |
There was a problem hiding this comment.
These tests can be merged in a single function, and use pytest.mark.parametrize to test both storages.
Also, in the end, I believe that if you add your test env in the list of envs as mentioned above, you can rely on previous written tests that should already check equality, and this file can be removed
| values = np.pad(values, ((0, pad), (0, 0))) | ||
| dtype = pa.list_(pa.from_numpy_dtype(space.dtype), list_size=values.shape[1]) | ||
| return pa.FixedSizeListArray.from_arrays(values.reshape(-1), type=dtype) | ||
| if values.shape == (4, 84, 84) and values.dtype == np.uint8: # check for image observation (4 stacked greyscale images) |
There was a problem hiding this comment.
we should not constrain ourselves with 84 x 84 images, but we should accept any size.
Same reasoning for the 4 dimension.
I think the discriminant here is to have Box, with at least 2 dim, and with type uint8. Then you can put a logging.warn saying you are considering it as image, and if this is not intended to disable it via a flag "image_observation" which is defaulted to None. Then you can compute the value of the flag during init (and warn just once). I will clarify later in our meeting.
| if values.type == pa.binary(): # check for binary data (JPEG) | ||
| jpeg_images = [] | ||
| for jpeg_bytes in values: | ||
| image = Image.open(io.BytesIO(jpeg_bytes)).convert("L") # decode JPEG and convert to greyscale |
There was a problem hiding this comment.
we should work with non-grayscale images as well.
Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
Fixes # (issue), Depends on # (pull request)
Type of change
Screenshots
Checklist:
pre-commitchecks withpre-commit run --all-files(seeCONTRIBUTING.mdinstructions to set it up)pytest -vand no errors are present.pytest -vhas generated that are related to my code to the best of my knowledge.