Merge pull request #13 from bioscan-ml/enh_arbitrary-modalities

ENH: Add support for arbitrary modalities

Author: scottclowe

README.rst

@@ -157,6 +157,27 @@ or ``"dna"``:
 
    dataset = BIOSCAN5M(root="~/Datasets/bioscan/", modality="dna")
 
+Additionally, any column names from the metadata can be used as input modalities.
+For example, to load the latitude and longitude coordinates as inputs:
+
+.. code-block:: python
+
+   dataset = BIOSCAN5M(root="~/Datasets/bioscan/", modality=("coord-lat", "coord-lon"))
+
+or to load the size of the insect (in pixels) in addition to the DNA barcode:
+
+.. code-block:: python
+
+   dataset = BIOSCAN5M(
+       root="~/Datasets/bioscan/", modality=("dna", "image_measurement_value")
+   )
+
+Multiple modalities can be selected by passing a list of column names.
+Each item in the dataset will have the inputs in the same order as specified in the ``modality`` argument.
+
+All samples have an image and a DNA barcode, but other fields may be incomplete.
+Any missing values will be replaced with NaN.
+
 
 Target selection
 ~~~~~~~~~~~~~~~~
@@ -239,14 +260,6 @@ The dataset class supports the use of data transforms for the image and DNA barc
    )
 
 
-Size and geolocation metadata
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-The BIOSCAN-5M dataset also contains insect size and geolocation metadata.
-Loading this metadata is not yet supported by the `BIOSCAN5M <BS5M-class_>`_ pytorch dataset class.
-In the meantime, users of the dataset are welcome to explore this metadata themselves.
-
-
 Other resources
 ---------------
 

bioscan_dataset/bioscan1m.py

@@ -237,7 +237,10 @@ class BIOSCAN1M(VisionDataset):
 
     modality : str or Iterable[str], default=("image", "dna")
         Which data modalities to use. One of, or a list of:
-        ``"image"``, ``"dna"``.
+        ``"image"``, ``"dna"``, or any column name in the metadata CSV file.
+
+        .. versionchanged:: 1.1.0
+            Added support for arbitrary modalities.
 
     reduce_repeated_barcodes : bool, default=False
         Whether to reduce the dataset to only one sample per barcode.
@@ -426,6 +429,13 @@ def __getitem__(self, index: int) -> Tuple[Any, ...]:
             The DNA barcode, if the ``"dna"`` modality is requested, optionally
             transformed by the ``dna_transform`` pipeline.
 
+        *modalities : Any
+            Any other modalities requested, as specified in the ``modality`` parameter.
+            The data is extracted from the appropriate column in the metadata TSV file,
+            without any transformations.
+
+            .. versionadded:: 1.1.0
+
         target : int or Tuple[int, ...] or str or Tuple[str, ...] or None
             The target(s), optionally transformed by the ``target_transform`` pipeline.
             If ``target_format="index"``, the target(s) will be returned as integer
@@ -446,6 +456,8 @@ def __getitem__(self, index: int) -> Tuple[Any, ...]:
                 X = sample["nucraw"]
                 if self.dna_transform is not None:
                     X = self.dna_transform(X)
+            elif modality in self.metadata.columns:
+                X = sample[modality]
             else:
                 raise ValueError(f"Unfamiliar modality: {modality}")
             values.append(X)
@@ -499,7 +511,7 @@ def _check_exists(self, verbose=0) -> bool:
 
     def _load_metadata(self) -> pandas.DataFrame:
         r"""
-        Load metadata from CSV file and prepare it for training.
+        Load metadata from TSV file and prepare it for training.
         """
         self.metadata = load_metadata(
             self.metadata_path,

bioscan_dataset/bioscan5m.py

@@ -62,6 +62,11 @@
     "species",
     "dna_bin",
     "dna_barcode",
+    "country",
+    "province_state",
+    "coord-lat",
+    "coord-lon",
+    "image_measurement_value",
     "split",
 ]
 
@@ -225,7 +230,10 @@ class BIOSCAN5M(VisionDataset):
 
     modality : str or Iterable[str], default=("image", "dna")
         Which data modalities to use. One of, or a list of:
-        ``"image"``, ``"dna"``.
+        ``"image"``, ``"dna"``, or any column name in the metadata CSV file.
+
+        .. versionchanged:: 1.1.0
+            Added support for arbitrary modalities.
 
     image_package : str, default="cropped_256"
         The package to load images from. One of:
@@ -507,6 +515,13 @@ def __getitem__(self, index: int) -> Tuple[Any, ...]:
             The DNA barcode, if the ``"dna"`` modality is requested, optionally
             transformed by the ``dna_transform`` pipeline.
 
+        *modalities : Any
+            Any other modalities requested, as specified in the ``modality`` parameter.
+            The data is extracted from the appropriate column in the metadata CSV file,
+            without any transformations.
+
+            .. versionadded:: 1.1.0
+
         target : int or Tuple[int, ...] or str or Tuple[str, ...] or None
             The target(s), optionally transformed by the ``target_transform`` pipeline.
             If ``target_format="index"``, the target(s) will be returned as integer
@@ -527,6 +542,8 @@ def __getitem__(self, index: int) -> Tuple[Any, ...]:
                 X = sample["dna_barcode"]
                 if self.dna_transform is not None:
                     X = self.dna_transform(X)
+            elif modality in self.metadata.columns:
+                X = sample[modality]
             else:
                 raise ValueError(f"Unfamiliar modality: {modality}")
             values.append(X)