edgefirst-validator 4.2.1__py3-none-any.whl

edgefirst/validator/datasets/cache.py

@@ -0,0 +1,296 @@
+"""
+Implementations for downloading and caching the dataset.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import glob
+from typing import TYPE_CHECKING, List, Tuple
+
+import tqdm
+import lmdb
+import numpy as np
+
+from edgefirst_client import Client, FileType, AnnotationType
+
+from edgefirst.validator.datasets import LMDBDatabase
+from edgefirst.validator.publishers.utils.logger import logger
+
+if TYPE_CHECKING:
+    from edgefirst.validator.evaluators import DatasetParameters
+    from edgefirst.validator.datasets import Dataset
+
+
+class StudioCache:
+    def __init__(
+        self,
+        parameters: DatasetParameters,
+        stages: List[Tuple[str, str]],
+        client: Client = None,
+        session_id: str = None,
+        val_group: str = "val"
+    ):
+        """
+        Communicate with EdgeFirst Studio for
+        fetching and caching the dataset and post the progress.
+
+        Parameters
+        ----------
+        parameters: DatasetParameters
+            This contains dataset parameters set from the command line.
+        stages: List[Tuple[str, str]]
+            This contains the stages that tracks each progress in Studio.
+            A stage contains ("stage identifier", "stage description").
+        client: Client
+            EdgeFirst Client object.
+        session_id: str
+            This is the validation session ID in EdgeFirst Studio for
+            posting validation metrics.
+        val_group: str
+            The dataset validation group set in EdgeFirst Studio.
+        """
+
+        self.parameters = parameters
+        self.stages = stages
+        self.client = client
+        self.val_session = None
+        self.val_group = val_group
+
+        if self.client is not None:
+            self.val_session = self.client.validation_session(session_id)
+
+    def complete_stage(self, stage: str, message: str):
+        """
+        Completes the stage on studio at the end of the task.
+
+        Parameters
+        ----------
+        stage: str
+            The stage identifier.
+        message: str
+            The message to project in the Studio GUI.
+        """
+        self.client.update_stage(
+            self.val_session.task.id,
+            stage=stage,
+            status="Running",
+            message=message,
+            percentage=100
+        )
+
+    def download(self, dataset: str):
+        """
+        Download the dataset from EdgeFirst Studio into the device.
+
+        Parameters
+        ----------
+        dataset: str
+            The path to directory to save the dataset.
+
+        Returns
+        -------
+        pl.DataFrame
+            This is the polars dataframe of the annotations.
+        """
+
+        if self.client is None:
+            raise ValueError(
+                "EdgeFirst Client needs to be defined to download the dataset.")
+
+        annotation_types = [AnnotationType.Box2d, AnnotationType.Mask]
+        os.makedirs(dataset, exist_ok=True)
+
+        dataset_id = self.val_session.dataset_id
+        annotation_id = self.val_session.annotation_set_id
+
+        # Download Images
+        with tqdm.tqdm(
+            total=0,
+            desc=f"Downloading Images from Dataset ID: ds-{dataset_id.value:x}"
+        ) as bar:
+            def progress(current, total):
+                if total != bar.total:
+                    bar.reset(total)
+                bar.update(current - bar.n)
+                self.client.update_stage(
+                    self.val_session.task.id,
+                    stage=self.stages[0][0],
+                    status="Running",
+                    message=self.stages[0][1],
+                    percentage=int(100 * current / total)
+                )
+
+            self.client.download_dataset(
+                dataset_id=dataset_id,
+                groups=[self.val_group],
+                types=[FileType.Image],
+                output=os.path.join(dataset, "images"),
+                progress=progress,
+            )
+
+        total_images = len(glob.glob(os.path.join(dataset, "images", "*"),
+                                     recursive=True))
+        logger(f"Downloaded a total of {total_images} images.", code="INFO")
+
+       # Download Annotations
+        with tqdm.tqdm(
+            total=0,
+            desc=f"Downloading Annotations from Annotation ID: as-{annotation_id.value:x}"
+        ) as bar:
+            def progress(current, total):
+                if total != bar.total:
+                    bar.reset(total)
+                bar.update(current - bar.n)
+                self.client.update_stage(
+                    self.val_session.task.id,
+                    stage=self.stages[1][0],
+                    status="Running",
+                    message=self.stages[1][1],
+                    percentage=int(100 * current / total)
+                )
+
+            dataframe = self.client.annotations_dataframe(
+                annotation_set_id=annotation_id,
+                groups=[self.val_group],
+                annotation_types=annotation_types,
+                progress=progress
+            )
+            dataframe.write_ipc(os.path.join(dataset, "dataset.arrow"))
+
+        logger(f"Downloaded a total of {dataframe.shape[0]} annotations.",
+               code="INFO")
+
+        return dataframe
+
+    def cache(self, dataset: Dataset,
+              cache: str = "cache/val.db") -> LMDBDatabase:
+        """
+        Cache the dataset provided into an LMDB file.
+
+        Parameters
+        ----------
+        dataset: Dataset
+            This can either be a DarkNet or EdgeFirst Dataset object.
+        cache: str
+            The path to the cache file to save the cache.
+
+        Returns
+        -------
+        LMDBDatabase
+            The instantiated cached dataset.
+        """
+
+        # Remove existing cache file.
+        if os.path.exists(cache):
+            os.remove(cache)
+        os.makedirs(os.path.dirname(cache), exist_ok=True)
+
+        dbenv = lmdb.open(
+            cache,
+            map_size=1024 ** 4,
+            max_dbs=10,
+            subdir=False,
+            lock=False
+        )
+
+        classes_db = dbenv.open_db(b'classes')  # Unique labels in the dataset.
+        # Name of each image in the dataset.
+        names_db = dbenv.open_db(b'names')
+        # All preprocessed images in the dataset.
+        images_db = dbenv.open_db(b'images')
+        # Images used for visualization.
+        visual_images_db = dbenv.open_db(b'visual')
+        boxes_db = dbenv.open_db(b'box2d')  # 2D box annotations.
+        # Integer labels for each boxes or masks.
+        labels_db = dbenv.open_db(b'labels')
+        masks_db = dbenv.open_db(b'masks')  # 2D masks annotations.
+
+        # Store the unique labels in the dataset.
+        with dbenv.begin(write=True) as txn:
+            txn.put(
+                b'classes',
+                json.dumps(self.parameters.labels).encode(),
+                db=classes_db
+            )
+
+        samples = tqdm.tqdm(dataset.collect_samples(), colour="green")
+        samples.set_description("Caching dataset")
+
+        for i, sample in enumerate(samples):
+            gt_instance = dataset.read_sample(sample)
+
+            with dbenv.begin(write=True) as txn:
+                image_path = gt_instance.image_path
+                image = gt_instance.image
+                visual_image = gt_instance.visual_image
+                shapes = gt_instance.shapes
+                image_shape = gt_instance.image_shape
+                labels = gt_instance.labels
+
+                if self.parameters.common.with_boxes:
+                    boxes = gt_instance.boxes
+                else:
+                    boxes = np.array([], dtype=np.float32)
+
+                if self.parameters.common.with_masks:
+                    masks = gt_instance.mask
+                else:
+                    masks = np.zeros(shapes[0], dtype=np.uint8)
+
+                name = os.path.basename(image_path)
+                txn.put(name.encode(), None, db=names_db)
+                txn.put(name.encode(), image.tobytes(), db=images_db)
+                # Place the current shape of the image.
+                txn.put(f'{name}/shape'.encode(),
+                        np.array(image.shape, dtype=np.int32).tobytes(),
+                        db=images_db)
+
+                if visual_image is not None:
+                    txn.put(
+                        name.encode(),
+                        visual_image.tobytes(),
+                        db=visual_images_db)
+                    # Place the current shape of the visualization image.
+                    txn.put(f'{name}/shape'.encode(),
+                            np.array(visual_image.shape,
+                                     dtype=np.int32).tobytes(),
+                            db=visual_images_db)
+
+                # Place the label transformation shapes after letterboxing.
+                txn.put(f'{name}/shapes'.encode(),
+                        json.dumps(shapes).encode(),
+                        db=images_db)
+                # Place the original shape of the image.
+                txn.put(f'{name}/im_shape'.encode(),
+                        np.array(image_shape, dtype=np.int32).tobytes(),
+                        db=images_db)
+                txn.put(name.encode(), boxes.tobytes(), db=boxes_db)
+                txn.put(name.encode(), labels.tobytes(), db=labels_db)
+                txn.put(name.encode(), masks.tobytes(), db=masks_db)
+                txn.put(f'{name}/mask_shape'.encode(),
+                        np.array(masks.shape, dtype=np.int32).tobytes(),
+                        db=masks_db)
+
+            if self.client is not None:
+                self.client.update_stage(
+                    self.val_session.task.id,
+                    stage=self.stages[2][0],
+                    status="Running",
+                    message=self.stages[2][1],
+                    percentage=int(100 * i / len(dataset))
+                )
+
+        dbenv.close()
+        if self.client is not None:
+            self.complete_stage(
+                stage=self.stages[2][0],
+                message=self.stages[2][1]
+            )
+
+        return LMDBDatabase(
+            source=cache,
+            parameters=self.parameters,
+            timer=dataset.timer
+        )

edgefirst/validator/datasets/core.py

@@ -0,0 +1,250 @@
+"""
+Common parent dataset implementations.
+"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING, Union, Tuple
+
+import numpy as np
+import tqdm
+
+from edgefirst.validator.datasets.utils.readers import read_image
+from edgefirst.validator.datasets.utils.transformations import (xcycwh2xyxy,
+                                                                xywh2xyxy,
+                                                                normalize,
+                                                                denormalize_polygon)
+
+if TYPE_CHECKING:
+    from edgefirst.validator.evaluators import DatasetParameters, TimerContext
+    from edgefirst_python import TensorImage  # type: ignore
+
+    from edgefirst.validator.datasets import (
+        SegmentationInstance, DetectionInstance, MultitaskInstance, Instance)
+
+
+class Dataset:
+    """
+    Abstract dataset class for providing template methods in the dataset.
+
+    Parameters
+    ----------
+    source: str
+        The path to the source dataset.
+    parameters: DatasetParameters
+        This contains dataset parameters set from the command line.
+    timer: TimerContext
+        A timer object for handling validation timings for the model.
+    info_dataset: dict
+        Contains information such as:
+
+            .. code-block:: python
+
+                {
+                    "classes": [list of unique labels],
+                    "validation":
+                    {
+                        "images: 'path to the images',
+                        "annotations": 'path to the annotations'
+                    }
+                }
+
+        *Note: the classes are optional and the path to the images
+        and annotations can be the same.*
+
+    Raises
+    ------
+    ValueError
+        Raised if the provided parameters in certain methods
+        does not conform to the specified data type.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        parameters: DatasetParameters,
+        timer: TimerContext,
+        info_dataset: dict = None
+    ):
+        self.source = source
+        self.parameters = parameters
+        self.timer = timer
+        self.info_dataset = info_dataset
+        self.samples = []
+
+        self.transformer = None
+        if self.parameters.box_format == 'xcycwh':
+            self.transformer = xcycwh2xyxy
+        elif self.parameters.box_format == 'xywh':
+            self.transformer = xywh2xyxy
+        else:
+            self.transformer = None
+
+        self.normalizer = None
+        self.denormalizer = None
+        if self.parameters.normalized:
+            if self.parameters.common.with_masks:
+                self.denormalizer = denormalize_polygon
+        else:
+            if self.parameters.common.with_boxes:
+                self.normalizer = normalize
+
+    def __len__(self) -> int:
+        """
+        Returns the number of samples in the dataset.
+        """
+        return len(self.samples)
+
+    def __iter__(self):
+        """
+        Reads all the samples in the dataset.
+
+        Yields
+        -------
+        Instance
+            Yields one sample of the ground truth
+            instance which contains information on the image
+            as a NumPy array, boxes, labels, and image path.
+        """
+        if self.parameters.silent:
+            samples = self.collect_samples()
+            for sample in samples:
+                yield self.read_sample(sample)
+        else:
+            samples = tqdm.tqdm(self.collect_samples(), colour="green")
+            samples.set_description("Validation Progress")
+            for sample in samples:
+                yield self.read_sample(sample)
+
+    def verify_dataset(self):
+        """Abstract Method"""
+        pass
+
+    def read_sample(self,
+                    sample: Union[list, Tuple[str, str], str]) -> Instance:
+        """
+        Reads one sample from the dataset.
+
+        Parameters
+        -----------
+        sample: Union[list, Tuple[str, str], str]
+            For EdgeFirstDatabase, this is a list. For Darknet datasets,
+            this is a Tuple[str, str] containing the path to the image
+            and annotations. For dataset cache, this is a string
+            as the image name.
+
+            A single dataset sample contains the indices
+            in the dataframe pointing to all the annotations
+            in the dataset for this sample.
+
+        Returns
+        -------
+        Instance
+            The ground truth instance objects contains the annotations
+            representing the ground truth of the image.
+        """
+        if self.parameters.common.with_boxes and self.parameters.common.with_masks:
+            return self.build_multitask_instance(sample)
+        elif self.parameters.common.with_boxes:
+            return self.build_detection_instance(sample)
+        elif self.parameters.common.with_masks:
+            return self.build_segmentation_instance(sample)
+        else:
+            raise ValueError(
+                "Could not determine model task as detection or segmentation.")
+
+    def load_image(
+        self,
+        image_path: str,
+        backend: str = "hal"
+    ) -> Union[TensorImage, np.ndarray]:
+        """
+        Load the image into memory using various libraries: "hal", "opencv",
+        or "pillow".
+
+        Parameters
+        ----------
+        image_path: str
+            The path to the image.
+        backend: str
+            Specify the backend library for resizing the image
+            from the options "hal", "opencv", "pillow".
+
+        Returns
+        -------
+        Union[edgefirst_python.TensorImage, np.ndarray]
+            TensorImage is returned when using "hal". Otherwise, a
+            NumPy array is returned.
+        """
+
+        if backend == "hal":
+            try:
+                import edgefirst_python  # type: ignore
+            except ImportError:
+                raise ImportError(
+                    "EdgeFirst HAL is needed to read the image.")
+            # Read the image.
+            return edgefirst_python.TensorImage.load(
+                image_path, fourcc=edgefirst_python.FourCC.RGBA)
+        elif backend == "opencv":
+            try:
+                import cv2  # type: ignore
+            except ImportError:
+                raise ImportError("OpenCV is needed to read the image.")
+
+            return cv2.imread(image_path)
+        else:
+            return read_image(image_path, rotate=True)
+
+    def image(self, sample: Union[tuple, list]):
+        """Abstract Method"""
+        raise NotImplementedError("Abstract Method")
+
+    def labels(self, sample: Union[tuple, list]) -> np.ndarray:
+        """
+        Fetch the labels of the specified sample.
+
+        Parameters
+        ----------
+        sample: Union[tuple, list]
+            A tuple containing the (image path, annotation path) or
+            a list of indices in the polars dataframe for the current sample.
+
+        Returns
+        -------
+        np.ndarray
+            The labels in the sample containing np.int32 elements.
+        """
+        return np.array([])
+
+    def boxes(self, index: int) -> np.ndarray:
+        """Abstract Method"""
+        raise NotImplementedError("Abstract Method")
+
+    def mask(self, index: int) -> np.ndarray:
+        """Abstract Method"""
+        raise NotImplementedError("Abstract Method")
+
+    def segments(self, index: int) -> np.ndarray:
+        """Abstract Method"""
+        raise NotImplementedError("Absract Method")
+
+    def name(self, index: int) -> str:
+        """Abstract Method"""
+        raise NotImplementedError("Abstract Method")
+
+    def collect_samples(self):
+        """Abstract Method"""
+        raise NotImplementedError("This is an abstract method.")
+
+    def build_detection_instance(self, sample: list) -> DetectionInstance:
+        """Abstract Method"""
+        pass
+
+    def build_segmentation_instance(
+            self, sample: list) -> SegmentationInstance:
+        """Abstract Method"""
+        pass
+
+    def build_multitask_instance(self, sample: list) -> MultitaskInstance:
+        """Abstract Method"""
+        pass