maite-datasets 0.0.1__py3-none-any.whl

maite_datasets/image_classification/_mnist.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+__all__ = []
+
+from pathlib import Path
+from typing import Any, Literal, Sequence, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from maite_datasets._base import BaseICDataset, DataLocation
+from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
+from maite_datasets._protocols import Transform
+
+MNISTClassStringMap = Literal[
+    "zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"
+]
+TMNISTClassMap = TypeVar(
+    "TMNISTClassMap", MNISTClassStringMap, int, list[MNISTClassStringMap], list[int]
+)
+CorruptionStringMap = Literal[
+    "identity",
+    "shot_noise",
+    "impulse_noise",
+    "glass_blur",
+    "motion_blur",
+    "shear",
+    "scale",
+    "rotate",
+    "brightness",
+    "translate",
+    "stripe",
+    "fog",
+    "spatter",
+    "dotted_line",
+    "zigzag",
+    "canny_edges",
+]
+
+
+class MNIST(BaseICDataset[NDArray[np.number[Any]]], BaseDatasetNumpyMixin):
+    """`MNIST <https://en.wikipedia.org/wiki/MNIST_database>`_ Dataset and `Corruptions <https://arxiv.org/abs/1906.02337>`_.
+
+    There are 15 different styles of corruptions. This class downloads differently depending on if you
+    need just the original dataset or if you need corruptions. If you need both a corrupt version and the
+    original version then choose `corruption="identity"` as this downloads all of the corrupt datasets and
+    provides the original as `identity`. If you just need the original, then using `corruption=None` will
+    download only the original dataset to save time and space.
+
+    Parameters
+    ----------
+    root : str or pathlib.Path
+        Root directory where the data should be downloaded to or the ``minst`` folder of the already downloaded data.
+    image_set : "train", "test" or "base", default "train"
+        If "base", returns all of the data to allow the user to create their own splits.
+    corruption : "identity", "shot_noise", "impulse_noise", "glass_blur", "motion_blur", \
+        "shear", "scale", "rotate", "brightness", "translate", "stripe", "fog", "spatter", \
+        "dotted_line", "zigzag", "canny_edges" or None, default None
+        Corruption to apply to the data.
+    transforms : Transform, Sequence[Transform] or None, default None
+        Transform(s) to apply to the data.
+    download : bool, default False
+        If True, downloads the dataset from the internet and puts it in root directory.
+        Class checks to see if data is already downloaded to ensure it does not create a duplicate download.
+    verbose : bool, default False
+        If True, outputs print statements.
+
+    Attributes
+    ----------
+    path : pathlib.Path
+        Location of the folder containing the data.
+    image_set : "train", "test" or "base"
+        The selected image set from the dataset.
+    index2label : dict[int, str]
+        Dictionary which translates from class integers to the associated class strings.
+    label2index : dict[str, int]
+        Dictionary which translates from class strings to the associated class integers.
+    metadata : DatasetMetadata
+        Typed dictionary containing dataset metadata, such as `id` which returns the dataset class name.
+    corruption : str or None
+        Corruption applied to the data.
+    transforms : Sequence[Transform]
+        The transforms to be applied to the data.
+    size : int
+        The size of the dataset.
+    
+    Note
+    ----
+    Data License: `CC BY 4.0 <https://creativecommons.org/licenses/by/4.0/>`_ for corruption dataset
+    """
+
+    _resources = [
+        DataLocation(
+            url="https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz",
+            filename="mnist.npz",
+            md5=False,
+            checksum="731c5ac602752760c8e48fbffcf8c3b850d9dc2a2aedcf2cc48468fc17b673d1",
+        ),
+        DataLocation(
+            url="https://zenodo.org/record/3239543/files/mnist_c.zip",
+            filename="mnist_c.zip",
+            md5=True,
+            checksum="4b34b33045869ee6d424616cd3a65da3",
+        ),
+    ]
+
+    index2label: dict[int, str] = {
+        0: "zero",
+        1: "one",
+        2: "two",
+        3: "three",
+        4: "four",
+        5: "five",
+        6: "six",
+        7: "seven",
+        8: "eight",
+        9: "nine",
+    }
+
+    def __init__(
+        self,
+        root: str | Path,
+        image_set: Literal["train", "test", "base"] = "train",
+        corruption: CorruptionStringMap | None = None,
+        transforms: Transform[NDArray[np.number[Any]]]
+        | Sequence[Transform[NDArray[np.number[Any]]]]
+        | None = None,
+        download: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        self.corruption = corruption
+        if self.corruption == "identity" and verbose:
+            print("Identity is not a corrupted dataset but the original MNIST dataset.")
+        self._resource_index = 0 if self.corruption is None else 1
+
+        super().__init__(
+            root,
+            image_set,
+            transforms,
+            download,
+            verbose,
+        )
+
+    def _load_data_inner(self) -> tuple[list[str], list[int], dict[str, Any]]:
+        """Function to load in the file paths for the data and labels from the correct data format"""
+        if self.corruption is None:
+            try:
+                file_path = self.path / self._resource.filename
+                self._loaded_data, labels = self._grab_data(file_path)
+            except FileNotFoundError:
+                self._loaded_data, labels = self._load_corruption()
+        else:
+            self._loaded_data, labels = self._load_corruption()
+
+        index_strings = np.arange(self._loaded_data.shape[0]).astype(str).tolist()
+        return index_strings, labels.tolist(), {}
+
+    def _load_corruption(self) -> tuple[NDArray[np.number[Any]], NDArray[np.uintp]]:
+        """Function to load in the file paths for the data and labels for the different corrupt data formats"""
+        corruption = self.corruption if self.corruption is not None else "identity"
+        base_path = self.path / "mnist_c" / corruption
+        if self.image_set == "base":
+            raw_data = []
+            raw_labels = []
+            for group in ["train", "test"]:
+                file_path = base_path / f"{group}_images.npy"
+                raw_data.append(self._grab_corruption_data(file_path))
+
+                label_path = base_path / f"{group}_labels.npy"
+                raw_labels.append(self._grab_corruption_data(label_path))
+
+            data = np.concatenate(raw_data, axis=0).transpose(0, 3, 1, 2)
+            labels = np.concatenate(raw_labels).astype(np.uintp)
+        else:
+            file_path = base_path / f"{self.image_set}_images.npy"
+            data = self._grab_corruption_data(file_path)
+            data = data.astype(np.float64).transpose(0, 3, 1, 2)
+
+            label_path = base_path / f"{self.image_set}_labels.npy"
+            labels = self._grab_corruption_data(label_path)
+            labels = labels.astype(np.uintp)
+
+        return data, labels
+
+    def _grab_data(
+        self, path: Path
+    ) -> tuple[NDArray[np.number[Any]], NDArray[np.uintp]]:
+        """Function to load in the data numpy array"""
+        with np.load(path, allow_pickle=True) as data_array:
+            if self.image_set == "base":
+                data = np.concatenate(
+                    [data_array["x_train"], data_array["x_test"]], axis=0
+                )
+                labels = np.concatenate(
+                    [data_array["y_train"], data_array["y_test"]], axis=0
+                ).astype(np.uintp)
+            else:
+                data, labels = (
+                    data_array[f"x_{self.image_set}"],
+                    data_array[f"y_{self.image_set}"].astype(np.uintp),
+                )
+            data = np.expand_dims(data, axis=1)
+        return data, labels
+
+    def _grab_corruption_data(self, path: Path) -> NDArray[np.number[Any]]:
+        """Function to load in the data numpy array for the previously chosen corrupt format"""
+        return np.load(path, allow_pickle=False)
+
+    def _read_file(self, path: str) -> NDArray[np.number[Any]]:
+        """
+        Function to grab the correct image from the loaded data.
+        Overwrite of the base `_read_file` because data is an all or nothing load.
+        """
+        index = int(path)
+        return self._loaded_data[index]

maite_datasets/image_classification/_ships.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+__all__ = []
+
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from maite_datasets._base import BaseICDataset, DataLocation
+from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
+from maite_datasets._protocols import Transform
+
+
+class Ships(BaseICDataset[NDArray[np.number[Any]]], BaseDatasetNumpyMixin):
+    """
+    A dataset that focuses on identifying ships from satellite images.
+
+    The dataset comes from kaggle,
+    `Ships in Satellite Imagery <https://www.kaggle.com/datasets/rhammell/ships-in-satellite-imagery>`_.
+    The images come from Planet satellite imagery when they gave
+    `open-access to a portion of their data <https://www.planet.com/pulse/open-california-rapideye-data/>`_.
+
+    There are 4000 80x80x3 (HWC) images of ships, sea, and land.
+    There are also 8 larger scene images similar to what would be operationally provided.
+
+    Parameters
+    ----------
+    root : str or pathlib.Path
+        Root directory where the data should be downloaded to or the ``ships`` folder of the already downloaded data.
+    transforms : Transform, Sequence[Transform] or None, default None
+        Transform(s) to apply to the data.
+    download : bool, default False
+        If True, downloads the dataset from the internet and puts it in root directory.
+        Class checks to see if data is already downloaded to ensure it does not create a duplicate download.
+    verbose : bool, default False
+        If True, outputs print statements.
+
+    Attributes
+    ----------
+    path : pathlib.Path
+        Location of the folder containing the data.
+    image_set : "base"
+        The base image set is the only available image set for the Ships dataset.
+    index2label : dict[int, str]
+        Dictionary which translates from class integers to the associated class strings.
+    label2index : dict[str, int]
+        Dictionary which translates from class strings to the associated class integers.
+    metadata : DatasetMetadata
+        Typed dictionary containing dataset metadata, such as `id` which returns the dataset class name.
+    transforms : Sequence[Transform]
+        The transforms to be applied to the data.
+    size : int
+        The size of the dataset.
+
+    Note
+    ----
+    Data License: `CC BY-SA 4.0 <https://creativecommons.org/licenses/by-sa/4.0/>`_
+    """
+
+    _resources = [
+        DataLocation(
+            url="https://zenodo.org/record/3611230/files/ships-in-satellite-imagery.zip",
+            filename="ships-in-satellite-imagery.zip",
+            md5=True,
+            checksum="b2e8a41ed029592b373bd72ee4b89f32",
+        ),
+    ]
+
+    index2label: dict[int, str] = {
+        0: "no ship",
+        1: "ship",
+    }
+
+    def __init__(
+        self,
+        root: str | Path,
+        transforms: Transform[NDArray[np.number[Any]]]
+        | Sequence[Transform[NDArray[np.number[Any]]]]
+        | None = None,
+        download: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            root,
+            "base",
+            transforms,
+            download,
+            verbose,
+        )
+        self._scenes: list[str] = self._load_scenes()
+        self._remove_extraneous_json_file()
+
+    def _remove_extraneous_json_file(self) -> None:
+        json_path = self.path / "shipsnet.json"
+        if json_path.exists():
+            json_path.unlink()
+
+    def _load_data_inner(self) -> tuple[list[str], list[int], dict[str, Any]]:
+        """Function to load in the file paths for the data and labels"""
+        file_data = {
+            "label": [],
+            "scene_id": [],
+            "longitude": [],
+            "latitude": [],
+            "path": [],
+        }
+        data_folder = sorted((self.path / "shipsnet").glob("*.png"))
+        if not data_folder:
+            raise FileNotFoundError
+
+        for entry in data_folder:
+            # Remove file extension and split by "_"
+            parts = entry.stem.split("__")
+            file_data["label"].append(int(parts[0]))
+            file_data["scene_id"].append(parts[1])
+            lat_lon = parts[2].split("_")
+            file_data["longitude"].append(float(lat_lon[0]))
+            file_data["latitude"].append(float(lat_lon[1]))
+            file_data["path"].append(entry)
+        data = file_data.pop("path")
+        labels = file_data.pop("label")
+        return data, labels, file_data
+
+    def _load_scenes(self) -> list[str]:
+        """Function to load in the file paths for the scene images"""
+        return sorted(str(entry) for entry in (self.path / "scenes").glob("*.png"))
+
+    def get_scene(self, index: int) -> NDArray[np.number[Any]]:
+        """
+        Get the desired satellite image (scene) by passing in the index of the desired file.
+
+        Args
+        ----
+        index : int
+            Value of the desired data point
+
+        Returns
+        -------
+        NDArray[np.number]
+            Scene image
+
+        Note
+        ----
+        The scene will be returned with the channel axis first.
+        """
+        scene = self._read_file(self._scenes[index])
+        np.moveaxis(scene, -1, 0)
+        return scene

maite_datasets/object_detection/__init__.py

@@ -0,0 +1,20 @@
+"""Module for MAITE compliant Object Detection datasets."""
+
+from maite_datasets.object_detection._antiuav import AntiUAVDetection
+from maite_datasets.object_detection._milco import MILCO
+from maite_datasets.object_detection._seadrone import SeaDrone
+from maite_datasets.object_detection._voc import VOCDetection
+
+__all__ = [
+    "AntiUAVDetection",
+    "MILCO",
+    "SeaDrone",
+    "VOCDetection",
+]
+
+import importlib.util
+
+if importlib.util.find_spec("torch") is not None:
+    from maite_datasets.object_detection._voc_torch import VOCDetectionTorch
+
+    __all__ += ["VOCDetectionTorch"]

maite_datasets/object_detection/_antiuav.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+__all__ = []
+
+from pathlib import Path
+from typing import Any, Literal, Sequence
+
+import numpy as np
+from defusedxml.ElementTree import parse
+from numpy.typing import NDArray
+
+from maite_datasets._base import BaseODDataset, DataLocation
+from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
+from maite_datasets._protocols import Transform
+
+
+class AntiUAVDetection(
+    BaseODDataset[NDArray[np.number[Any]], list[str], str], BaseDatasetNumpyMixin
+):
+    """
+    A UAV detection dataset focused on detecting UAVs in natural images against large variation in backgrounds.
+
+    The dataset comes from the paper
+    `Vision-based Anti-UAV Detection and Tracking <https://ieeexplore.ieee.org/document/9785379>`_
+    by Jie Zhao et. al. (2022).
+
+    The dataset is approximately 1.3 GB and can be found `here <https://github.com/wangdongdut/DUT-Anti-UAV>`_.
+    Images are collected against a variety of different backgrounds with a variety in the number and type of UAV.
+    Ground truth labels are provided for the train, validation and test set.
+    There are 35 different types of drones along with a variety in lighting conditions and weather conditions.
+
+    There are 10,000 images: 5200 images in the training set, 2200 images in the validation set,
+    and 2600 images in the test set.
+    The dataset only has a single UAV class with the focus being on identifying object location in the image.
+    Ground-truth bounding boxes are provided in (x0, y0, x1, y1) format.
+    The images come in a variety of sizes from 3744 x 5616 to 160 x 240.
+
+    Parameters
+    ----------
+    root : str or pathlib.Path
+        Root directory where the data should be downloaded to or
+        the ``antiuavdetection`` folder of the already downloaded data.
+    image_set: "train", "val", "test", or "base", default "train"
+        If "base", then the full dataset is selected (train, val and test).
+    transforms : Transform, Sequence[Transform] or None, default None
+        Transform(s) to apply to the data.
+    download : bool, default False
+        If True, downloads the dataset from the internet and puts it in root directory.
+        Class checks to see if data is already downloaded to ensure it does not create a duplicate download.
+    verbose : bool, default False
+        If True, outputs print statements.
+
+    Attributes
+    ----------
+    path : pathlib.Path
+        Location of the folder containing the data.
+    image_set : "train", "val", "test", or "base"
+        The selected image set from the dataset.
+    index2label : dict[int, str]
+        Dictionary which translates from class integers to the associated class strings.
+    label2index : dict[str, int]
+        Dictionary which translates from class strings to the associated class integers.
+    metadata : DatasetMetadata
+        Typed dictionary containing dataset metadata, such as `id` which returns the dataset class name.
+    transforms : Sequence[Transform]
+        The transforms to be applied to the data.
+    size : int
+        The size of the dataset.
+
+    Note
+    ----
+    Data License: `Apache 2.0 <https://www.apache.org/licenses/LICENSE-2.0.txt>`_
+    """
+
+    # Need to run the sha256 on the files and then store that
+    _resources = [
+        DataLocation(
+            url="https://drive.usercontent.google.com/download?id=1RVsSGPUKTdmoyoPTBTWwroyulLek1eTj&export=download&authuser=0&confirm=t&uuid=6bca4f94-a242-4bc2-9663-fb03cd94ef2c&at=APcmpox0--NroQ_3bqeTFaJxP7Pw%3A1746552902927",
+            filename="train.zip",
+            md5=False,
+            checksum="14f927290556df60e23cedfa80dffc10dc21e4a3b6843e150cfc49644376eece",
+        ),
+        DataLocation(
+            url="https://drive.usercontent.google.com/download?id=1333uEQfGuqTKslRkkeLSCxylh6AQ0X6n&export=download&authuser=0&confirm=t&uuid=c2ad2f01-aca8-4a85-96bb-b8ef6e40feea&at=APcmpozY-8bhk3nZSFaYbE8rq1Fi%3A1746551543297",
+            filename="val.zip",
+            md5=False,
+            checksum="238be0ceb3e7c5be6711ee3247e49df2750d52f91f54f5366c68bebac112ebf8",
+        ),
+        DataLocation(
+            url="https://drive.usercontent.google.com/download?id=1L1zeW1EMDLlXHClSDcCjl3rs_A6sVai0&export=download&authuser=0&confirm=t&uuid=5a1d7650-d8cd-4461-8354-7daf7292f06c&at=APcmpozLQC1CuP-n5_UX2JnP53Zo%3A1746551676177",
+            filename="test.zip",
+            md5=False,
+            checksum="a671989a01cff98c684aeb084e59b86f4152c50499d86152eb970a9fc7fb1cbe",
+        ),
+    ]
+
+    index2label: dict[int, str] = {
+        0: "unknown",
+        1: "UAV",
+    }
+
+    def __init__(
+        self,
+        root: str | Path,
+        image_set: Literal["train", "val", "test", "base"] = "train",
+        transforms: Transform[NDArray[np.number[Any]]]
+        | Sequence[Transform[NDArray[np.number[Any]]]]
+        | None = None,
+        download: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            root,
+            image_set,
+            transforms,
+            download,
+            verbose,
+        )
+
+    def _load_data(self) -> tuple[list[str], list[str], dict[str, list[Any]]]:
+        filepaths: list[str] = []
+        targets: list[str] = []
+        datum_metadata: dict[str, list[Any]] = {}
+
+        # If base, load all resources
+        if self.image_set == "base":
+            metadata_list: list[dict[str, Any]] = []
+
+            for resource in self._resources:
+                self._resource = resource
+                resource_filepaths, resource_targets, resource_metadata = (
+                    super()._load_data()
+                )
+                filepaths.extend(resource_filepaths)
+                targets.extend(resource_targets)
+                metadata_list.append(resource_metadata)
+
+            # Combine metadata
+            for data_dict in metadata_list:
+                for key, val in data_dict.items():
+                    str_key = str(key)  # Ensure key is string
+                    if str_key not in datum_metadata:
+                        datum_metadata[str_key] = []
+                    datum_metadata[str_key].extend(val)
+
+        else:
+            # Grab only the desired data
+            for resource in self._resources:
+                if self.image_set in resource.filename:
+                    self._resource = resource
+                    resource_filepaths, resource_targets, resource_metadata = (
+                        super()._load_data()
+                    )
+                    filepaths.extend(resource_filepaths)
+                    targets.extend(resource_targets)
+                    datum_metadata.update(resource_metadata)
+
+        return filepaths, targets, datum_metadata
+
+    def _load_data_inner(self) -> tuple[list[str], list[str], dict[str, Any]]:
+        resource_name = self._resource.filename[:-4]
+        base_dir = self.path / resource_name
+        data_folder = sorted((base_dir / "img").glob("*.jpg"))
+        if not data_folder:
+            raise FileNotFoundError
+
+        file_data = {
+            "image_id": [f"{resource_name}_{entry.name}" for entry in data_folder]
+        }
+        data = [str(entry) for entry in data_folder]
+        annotations = sorted(str(entry) for entry in (base_dir / "xml").glob("*.xml"))
+
+        return data, annotations, file_data
+
+    def _read_annotations(
+        self, annotation: str
+    ) -> tuple[list[list[float]], list[int], dict[str, Any]]:
+        """Function for extracting the info for the label and boxes"""
+        boxes: list[list[float]] = []
+        labels = []
+        root = parse(annotation).getroot()
+        if root is None:
+            raise ValueError(f"Unable to parse {annotation}")
+        additional_meta: dict[str, Any] = {
+            "image_width": int(root.findtext("size/width", default="-1")),
+            "image_height": int(root.findtext("size/height", default="-1")),
+            "image_depth": int(root.findtext("size/depth", default="-1")),
+        }
+        for obj in root.findall("object"):
+            labels.append(1 if obj.findtext("name", default="") == "UAV" else 0)
+            boxes.append(
+                [
+                    float(obj.findtext("bndbox/xmin", default="0")),
+                    float(obj.findtext("bndbox/ymin", default="0")),
+                    float(obj.findtext("bndbox/xmax", default="0")),
+                    float(obj.findtext("bndbox/ymax", default="0")),
+                ]
+            )
+
+        return boxes, labels, additional_meta