Add optional metadata output (DatasetSpecs fields) to DataLoader

[Copilot]

Extends ZarrDataset to optionally return complete DatasetSpecs metadata alongside samples for tracking sample provenance during training/inference workflows.

Changes

• Added return_metadata parameter to ZarrDataset: When enabled, appends metadata dict as last element in output tuple containing all standard DatasetSpecs fields for each modality: filenames (string representation of source), source_axes, axes, data_group (resolution level), and roi (formatted as human-readable string)

• Created zarrdataset_collate_fn: Custom collate function required for PyTorch DataLoader to properly batch non-numeric metadata separately from tensors. Detects metadata by checking for expected DatasetSpecs keys.

• Added format_roi() helper function: Converts ROI slices to human-readable string format (start_coords):(lengths), consistent with the ROI parsing convention (e.g., (10,20,0):(100,100,-1) or : for full selection)

• Maintains backward compatibility: Metadata placed last in output tuple (after worker_id, positions, arrays). Existing code unaffected when flag disabled.

Usage

import zarrdataset as zds
from torch.utils.data import DataLoader

ds = zds.ZarrDataset(
    dataset_specs=specs,
    return_metadata=True
)

loader = DataLoader(
    ds,
    batch_size=16,
    collate_fn=zds.zarrdataset_collate_fn,
    worker_init_fn=zds.zarrdataset_worker_init_fn
)

for *tensors, metadata_list in loader:
    for metadata in metadata_list:
        # metadata is dict keyed by modality (e.g., 'images', 'labels')
        for modality, mod_meta in metadata.items():
            # mod_meta contains: filenames, source_axes, axes, data_group, roi
            print(f"{modality}: {mod_meta['filenames']}, roi: {mod_meta['roi']}")

Output Order

• When combined with other flags: worker_id → positions → arrays → metadata
• Metadata always last when enabled

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Add Optional filename and data_scale Outputs to DataLoader</issue_title>
<issue_description># Add Optional filename and dataset_scale Outputs to DataLoader

Summary

Propose extending the current DataLoader behavior to optionally return filename and data_scale alongside each data sample during training or inference.

Motivation

In some workflows (e.g., inference, debugging, or evaluation), it's valuable to track which file a given sample originated from, as well as the dataset resolution or zoom level (data_scale). Currently, this metadata is not returned from the DataLoader, making it harder to:

• Associate predictions with specific files or scales
• Log or visualize results contextually
• Perform file- or scale-specific postprocessing or filtering

Proposed Change

Add a flag (e.g. return_metadata: bool = False) to the dataset/dataloader spec that, when enabled, changes the output format from: x to: x, metadata, where metadata is a dictionary or tuple containing:

{
  "filename": str,
  "dataset_scale": Optional[int]  # or float if appropriate
}

Alternatively, the return could be a named tuple for consistency.

Usage Example

for x, meta in dataloader:
    print(meta["filename"], meta["dataset_scale"])

Impact

This change would be fully backward-compatible. By default, the DataLoader would continue returning only x. When return_metadata=True, the additional values are returned.

Open Questions

• Should metadata be a dictionary, namedtuple, or part of a structured dataclass?
• Should the dataset_scale be inferred, passed in, or tied to the sample's spec?
• Are there training scenarios where metadata should be excluded for performance?</issue_description>

Comments on the Issue (you are @copilot in this section)

@fercer Sorry for the late follow up.

Yes, I think making all relevant metadata generated with the sampling process is valuable and would require some sort of structure.
Currently that is being yielded as additional outputs with return_positions and return_worker_id arguments.
The reason is that those values can be converted into pytroch tensors and handled by the default collate function.

For returning non-numeric values, like the filenames or other metadata, a custom collate function would be necesarry.
That function could merge all the metadata into a single dictionary as you suggested.
The collate function would be shipped within ZarrDataset, just as the worker_init_fn that should be passed to the DataLoader.

I'll start looking into how can the metadata be collated and delivered, any other ideas are welcome!</comment_new>

• Fixes Add Optional filename and data_scale Outputs to DataLoader #24

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[fercer]

This really should return the standard metadata from the DatasetSpecs class

[Copilot]

Updated the docstring to clarify that return_metadata returns all standard DatasetSpecs fields (filenames, source_axes, axes, data_group, roi) for each modality. (3883974)

[fercer]

Should roi be also converted into a string. Maybe with the format used to define ROIs within the same context.

[Copilot]

Added format_roi() helper function that converts ROI slices to string format (start_coords):(lengths), consistent with the format used in the metadata parsing context. For example: (0,10,5):(100,50,3) or : for full slices. (3883974)