partition_dataset_classes should output the partitioned labels too

[JHancox]

Is your feature request related to a problem? Please describe.
If you use partion_dataset_classes, it only returns the partitioned images and not the labels. This means that you have to, instead, pass it a numeric range instead of the images and then use that to get the labels and images. This additional step is not obvious from the documentation, adds additional complexity and seems unneccessary.

Describe the solution you'd like
The function should return the labels as a second return parameter

Describe alternatives you've considered
At the very least, improving the documentation to explain the current usage when you have class labels
If compatibility is an issue, then at least have a non-default parameter, such as return_labels=True

Additional context
Noticed this is the Monai-End-to-End notebook and figured not many folks would work this out intuitively

[wyli]

hi @JHancox the first argument also accepts a sequence of dictionaries, for example:

>>> from monai.data.utils import partition_dataset_classes
>>> data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
>>> classes = [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3]
>>> data_dicts = [{'image': x, 'label': y} for x, y in zip(data, classes)]
>>> partition_dataset_classes(data_dicts, classes, shuffle=False, ratios=[2,1])
[[{'image': 2, 'label': 0}, {'image': 8, 'label': 0}, {'image': 4, 'label': 1}, {'image': 1, 'label': 2}, {'image': 3, 'label': 2}, {'image': 6, 'label': 2}, {'image': 5, 'label': 3}, {'image': 11, 'label': 3}, {'image': 12, 'label': 3}], [{'image': 10, 'label': 0}, {'image': 13, 'label': 1}, {'image': 7, 'label': 2}, {'image': 9, 'label': 2}, {'image': 14, 'label': 3}]]

does this work for you?