mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-03-13 04:45:32 +08:00
65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
"""InsDet Scenes dataset for multi-object retrieval benchmark."""
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from benchmarks.base import BaseDataset
|
|
from data_loading.loader import load_val_dataset
|
|
|
|
|
|
class InsDetScenesDataset(BaseDataset):
|
|
"""InsDet-FULL/Scenes dataset with easy/hard splits.
|
|
|
|
This dataset provides scene images with object annotations from the
|
|
Instance Detection (InsDet) dataset, supporting easy and hard splits.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
scenes_dir: Path | str,
|
|
split: str = "easy",
|
|
):
|
|
"""Initialize InsDet Scenes dataset.
|
|
|
|
Args:
|
|
scenes_dir: Path to the InsDet-FULL/Scenes directory.
|
|
split: Scene split to use ('easy' or 'hard').
|
|
"""
|
|
self.scenes_dir = Path(scenes_dir)
|
|
self.split = split
|
|
self._dataset = load_val_dataset(self.scenes_dir, split)
|
|
|
|
def get_train_split(self) -> Any:
|
|
"""Get training split (same as test for this dataset).
|
|
|
|
Returns:
|
|
HuggingFace Dataset for training.
|
|
"""
|
|
return self._dataset
|
|
|
|
def get_test_split(self) -> Any:
|
|
"""Get test/evaluation split.
|
|
|
|
Returns:
|
|
HuggingFace Dataset for testing.
|
|
"""
|
|
return self._dataset
|
|
|
|
def __len__(self) -> int:
|
|
"""Get dataset length."""
|
|
return len(self._dataset)
|
|
|
|
def __getitem__(self, idx: int) -> dict[str, Any]:
|
|
"""Get a single item from the dataset.
|
|
|
|
Args:
|
|
idx: Index of the item.
|
|
|
|
Returns:
|
|
Dictionary containing:
|
|
- image: PIL Image
|
|
- image_id: Scene identifier
|
|
- objects: dict with bbox, category, area, id
|
|
"""
|
|
return self._dataset[idx]
|