utils.get_stored_demos and dataset_generator.py don't agree on mask image format, results in corrupted handles for stored demos

[Ben Eisner]

I've been using the dataset_generator.py script to generate demos for RLBench tasks, which is great.

However, when I try to load them, I find that the segmentation mask handles are different when reading from a stored dataset (i.e. by calling utils.get_stored_demos) compared to loading them online.

I've tracked this down to the fact that the dataset_generator.py script converts 3-channel segmentation masks into the [0, 255] range and converted into an int before storing as pngs: https://github.com/stepjam/RLBench/blob/master/tools/dataset_generator.py#L96

However, when the demos are read in utils.get_stored_demos, the PNGs generated by this process are read here: https://github.com/stepjam/RLBench/blob/master/rlbench/utils.py#L285, which converts them from 3-channel masks to 1-channel masks using this function: https://github.com/stepjam/RLBench/blob/master/rlbench/backend/utils.py#L223

As you can see, the mask is multiplied by 255 again, and converted into int. However, when it is read in, it is already on the range [0, 255], and of type uint8. Therefore multiplying by 255 causes some overflow, and i can't think of a way to invert the resulting overflow.

I've currently modified rgb_handles_to_mask to remove this conversion for my purposes, but im not sure if there are other use cases where that conversion should be handled. Is dataset_generator.py the only producer of PNGs which would be consumed this way?