diff --git a/kenning/datasets/lindenthal_camera_traps_dataset.py b/kenning/datasets/lindenthal_camera_traps_dataset.py index fbf92d82..32d21836 100644 --- a/kenning/datasets/lindenthal_camera_traps_dataset.py +++ b/kenning/datasets/lindenthal_camera_traps_dataset.py @@ -25,7 +25,12 @@ SegmObject, ) from kenning.utils.logger import KLogger, LoggerProgressBar -from kenning.utils.resource_manager import ResourceManager, Resources +from kenning.utils.resource_manager import ( + ResourceManager, + Resources, + ResourceURI, + extract_zip, +) try: import rosbags @@ -322,6 +327,10 @@ class LindenthalCameraTrapsDataset(ObjectDetectionSegmentationDataset): } ) + demo_dataset_uri = ResourceURI( + "kenning:///datasets/lindenthal-dataset-sample.zip" + ) + arguments_structure = { "image_width": { "description": "Width of the input images", @@ -333,6 +342,11 @@ class LindenthalCameraTrapsDataset(ObjectDetectionSegmentationDataset): "type": int, "default": 0, }, + "use_demonstration_dataset": { + "description": "If set to True, then instead of downloading a full Lindenthal dataset, a smaller, demonstration dataset is downloaded", # noqa: E501 + "type": bool, + "default": False, + }, } def __init__( @@ -353,9 +367,11 @@ def __init__( min_iou: float = 0.5, max_preds: int = 100, augment: bool = False, + use_demonstration_dataset: bool = False, ): self.num_classes = 4 self.augment = augment + self.use_demonstration_dataset = use_demonstration_dataset super().__init__( root=root, @@ -375,8 +391,21 @@ def __init__( max_preds=max_preds, ) - def download_dataset_fun(self): - self.root.mkdir(parents=True, exist_ok=True) + def download_original_dataset(self): + """ + Downloads the Lindenthal dataset. + + The dataset consists of ROS 1 rosbag files that contain data + from various sensors - camera, depth camera, infrared camera, ... + + This method downloads the dataset, reads rosbag files, extracts + the color and infrared images from it, and loads annotations for it. + + NOTE: At download, it reaches a size of around 220 GB. + The downloaded file is later removed, but it is recommended + to prepare a storage or configure a different one with + `KENNING_CACHE_DIR` + """ topics = [ "/device_0/sensor_0/Infrared_1/image/data", "/device_0/sensor_1/Color_0/image/data", @@ -404,6 +433,28 @@ def download_dataset_fun(self): (self.root / f).unlink() shutil.rmtree((self.root / "lindenthal-camera-traps")) + def download_demonstration_dataset(self): + """ + Downloads sample sequences from Lindenthal dataset. + + The sample sequences are of size around 1GB, providing + ready-to-use images and annotations (without the need + for ROS 1 rosbag unpacking). + + It is mostly for testing/demonstration purposes, it + is not recommended for training due to size. + + It should be used with `split_fraction_test` set to 1.0. + """ + extract_zip(self.root, self.demo_dataset_uri) + + def download_dataset_fun(self): + self.root.mkdir(parents=True, exist_ok=True) + if self.use_demonstration_dataset: + self.download_demonstration_dataset() + else: + self.download_original_dataset() + def __next__(self): X, y = super().__next__() if self.augment: