edgefirst-validator 4.1.12__py3-none-any.whl

edgefirst/validator/datasets/cache.py

@@ -0,0 +1,258 @@
+"""
+Implementations for downloading and caching the dataset.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import glob
+from typing import TYPE_CHECKING
+
+import tqdm
+import lmdb
+import numpy as np
+
+from edgefirst_client import Client, FileType, AnnotationType
+
+from edgefirst.validator.datasets import EdgeFirstDatabase
+from edgefirst.validator.publishers.utils.logger import logger
+
+if TYPE_CHECKING:
+    from edgefirst.validator.evaluators import DatasetParameters
+
+
+class StudioCache:
+    def __init__(
+        self,
+        parameters: DatasetParameters,
+        client: Client,
+        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.
+        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.client = client
+        self.edgefirst_dataset = None
+
+        self.val_session = self.client.validation_session(session_id)
+        self.val_task_id = self.val_session.task.id
+        train_session_id = self.val_session.training_session_id
+        self.train_session = self.client.training_session(train_session_id)
+
+        self.val_group = val_group
+
+    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.
+        """
+        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="fetch_img",
+                    status="Running",
+                    description="Downloading Images",
+                    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="fetch_as",
+                    status="Running",
+                    description="Downloading Annotations",
+                    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: str, cache: str = "cache/val.db"):
+        """
+        Cache the dataset provided into an LMDB file.
+
+        Parameters
+        ----------
+        dataset: str
+            The path to the dataset directory.
+        cache: str
+            The path to the cache file to save the cache.
+        """
+        group_name = os.path.basename(cache)
+        logger(f"Caching dataset to {group_name}", code="INFO")
+
+        os.makedirs(os.path.dirname(cache), exist_ok=True)
+
+        self.edgefirst_dataset = EdgeFirstDatabase(
+            source=dataset,
+            parameters=self.parameters,
+        )
+        ds_iterator = tqdm.tqdm(self.edgefirst_dataset)
+
+        # Remove existing cache file.
+        if os.path.exists(cache):
+            os.remove(cache)
+
+        dbenv = lmdb.open(
+            cache,
+            map_size=1024 ** 4,
+            max_dbs=10,
+            subdir=False,
+            lock=False
+        )
+
+        names_db = dbenv.open_db(b'names')
+        camera_db = dbenv.open_db(b'camera')
+        labels_db = dbenv.open_db(b'labels')
+
+        boxes_db = None
+        if self.parameters.common.with_boxes:
+            boxes_db = dbenv.open_db(b'box2d')
+
+        masks_db = None
+        if self.parameters.common.with_masks:
+            masks_db = dbenv.open_db(b'masks')
+
+        with dbenv.begin(write=True) as txn:
+            txn.put(
+                b'labels',
+                json.dumps(self.parameters.labels).encode(),
+                db=labels_db
+            )
+
+        for i, instance in enumerate(ds_iterator):
+            with dbenv.begin(write=True) as txn:
+                boxes = None
+                masks = None
+
+                if self.parameters.common.with_boxes and self.parameters.common.with_masks:
+                    name, image_data, boxes, (masks, masks_shape) = instance
+                elif self.parameters.common.with_boxes:
+                    name, image_data, boxes = instance
+                elif self.parameters.common.with_masks:
+                    name, image_data, (masks, masks_shape) = instance
+
+                image, shapes, ratio, image_shape = image_data
+
+                if boxes is not None:
+                    txn.put(name.encode(), boxes.tobytes(), db=boxes_db)
+
+                if masks is not None:
+                    masks = masks.astype(np.uint8).tobytes()
+                    txn.put(name.encode(), masks, db=masks_db)
+
+                txn.put(name.encode(), None, db=names_db)
+
+                if image is not None:
+                    txn.put(name.encode(), image.tobytes(), db=camera_db)
+
+                if shapes is not None:
+                    txn.put(f'{name}/shapes'.encode(),
+                            json.dumps(shapes).encode(),
+                            db=camera_db)
+
+                if ratio is not None:
+                    txn.put(f'{name}/ratio'.encode(),
+                            np.array([ratio], dtype=np.float32).tobytes(),
+                            db=camera_db)
+
+                # These are the original image dimensions.
+                if image_shape is not None:
+                    txn.put(f'{name}/shape'.encode(),
+                            np.array(image_shape, dtype=np.int32).tobytes(),
+                            db=camera_db)
+
+                # These are the shape of the mask.
+                if masks_shape is not None:
+                    txn.put(f'{name}/mask_shape'.encode(),
+                            np.array(masks_shape, dtype=np.int32).tobytes(),
+                            db=masks_db)
+
+            self.client.update_stage(
+                self.val_session.task.id,
+                stage="cache",
+                status="Running",
+                description=f"Caching: {group_name}",
+                percentage=int(100 * i / len(ds_iterator))
+            )
+
+        # Close the database.
+        dbenv.close()
+
+        self.client.update_stage(
+            self.val_session.task.id,
+            stage="cache",
+            status="Running",
+            description=f"Caching: {group_name}",
+            percentage=100
+        )

edgefirst/validator/datasets/core.py

@@ -0,0 +1,245 @@
+"""
+Common parent dataset implementations.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+from PIL import Image
+
+from edgefirst.validator.datasets.instance import Instance
+from edgefirst.validator.datasets.utils.transformations import (yolo2xyxy,
+                                                                xywh2xyxy,
+                                                                normalize,
+                                                                denormalize_polygon)
+
+if TYPE_CHECKING:
+    from deepview.datasets.generators.detection import BaseObjectDetectionGenerator  # type: ignore
+    from deepview.datasets.readers.darknet import DarknetDetectionReader  # type: ignore
+    from edgefirst.validator.evaluators import DatasetParameters
+
+
+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.
+
+    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
+    ):
+        self.source = source
+        self.parameters = parameters
+
+        self.transformer = None
+        if self.parameters.box_format == 'yolo':
+            self.transformer = yolo2xyxy
+        elif self.parameters.box_format == 'coco':
+            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
+
+        self.read_timings = list()
+        self.load_timings = list()
+
+    def build_dataset(self):
+        """Abstract Method"""
+        raise NotImplementedError("This is an abstract method.")
+
+    def read_sample(self, instance):
+        """Abstract Method"""
+        raise NotImplementedError("This is an abstract method.")
+
+    def read_all_samples(
+        self,
+        info: str = "Validation Progress",
+        silent: bool = False
+    ):
+        """
+        Reads all the samples in the dataset.
+
+        Parameters
+        ----------
+        info: str
+            The description of why image instances are being read.
+            By default it is to run validation, hence "Validation Progress".
+        silent: bool
+            If set to true, prevent validation logging.
+
+        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 silent:
+            samples = self.build_dataset()
+            for sample in samples:
+                yield self.read_sample(sample)
+        else:
+            try:
+                from tqdm import tqdm
+            except ImportError:
+                pass
+
+            try:
+                samples = tqdm(self.build_dataset(), colour="green")
+                samples.set_description(info)
+                for sample in samples:
+                    yield self.read_sample(sample)
+            except NameError:
+                samples = self.build_dataset()
+                num_samples = len(samples)
+                for index in range(num_samples):
+                    print("\t - [INFO]: Computing metrics for image: " +
+                          "%i of %i [%2.f %s]" %
+                          (index + 1,
+                           num_samples,
+                           100 * ((index + 1) / float(num_samples)),
+                           '%'), end='\r')
+                    yield self.read_sample(samples[index])
+
+    def timings(self):
+        """
+        Returns a summary of all the timings:
+        (mean, avg, max) of the preprocessing time.
+
+        Returns
+        -------
+            timings in ms: dict
+
+            .. code-block:: python
+
+                {
+                 'min_input_time': minimum time to load an image,
+                 'max_input_time': maximum time to load an image,
+                 'avg_input_time': average time to load an image,
+                }
+        """
+        return {
+            'min_input_time': (np.min(self.load_timings)
+                               if len(self.load_timings) else 0),
+            'max_input_time': (np.max(self.load_timings)
+                               if len(self.load_timings) else 0),
+            'avg_input_time': (np.mean(self.load_timings)
+                               if len(self.load_timings) else 0),
+        }
+
+
+class BaseDataset(Dataset):
+    """
+    This class utilizes deepview-datasets methods for iterating through
+    the images and annotations.
+
+    Parameters
+    ----------
+    source: str
+        The path to the dataset.
+    iterator: Iterator
+        Object in deepview-datasets for iterating through
+        the images or annotations. This can either be a generator if
+        a YAML file was passed, or a Reader if a directory was passed.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        iterator: Union[BaseObjectDetectionGenerator, DarknetDetectionReader]
+    ):
+        super(BaseDataset, self).__init__(source)
+        self.iterator = iterator
+
+        if isinstance(self.iterator, BaseObjectDetectionGenerator):
+            self.storage = self.iterator.reader.storage
+            self.labels = self.iterator.reader.classes
+        else:
+            self.storage = self.iterator.storage
+            self.labels = self.iterator.classes
+
+    def build_dataset(
+            self) -> Union[BaseObjectDetectionGenerator, DarknetDetectionReader]:
+        """
+        Returns the iterator object which already contains all the images
+        and annotations read in the dataset.
+
+        Returns
+        -------
+        Union[BaseObjectDetectionGenerator, DarknetDetectionReader]
+            BaseObjectDetectionGenerator
+                A generator if a YAML file was passed.
+            DarknetDetectionReader
+                Reader if a directory was passed.
+        """
+        return self.iterator
+
+    def read_sample(self, sample: tuple) -> Instance:
+        """
+        Returns the ground truth instance object which is needed to be read
+        by validator.
+
+        Parameters
+        ----------
+        sample: tuple
+            This contains the (image, boxes) in one sample.
+
+        Returns
+        -------
+        Instance
+            An object that contains the image, boxes, labels, etc.
+        """
+        from edgefirst.validator.datasets.utils.transformations import yolo2xyxy
+
+        image, boxes = sample
+        if len(image.shape) < 2:
+            image = Image.open(io.BytesIO(image)).convert('RGB')
+            image = np.asarray(image, dtype=np.uint8)
+        height, width, _ = image.shape
+
+        if isinstance(self.iterator, BaseObjectDetectionGenerator):
+            image_path = self.iterator.reader.get_instance_id()
+        else:
+            image_path = self.iterator.get_instance_id()
+
+        # Add file extension to allow image saving in disk.
+        if os.path.splitext(image_path)[-1] == "":
+            image_path += ".png"
+
+        instance = Instance(image_path)
+        instance.height = height
+        instance.width = width
+        instance.image = image
+
+        boxes = boxes[np.sum(boxes, axis=-1) != 0]
+        instance.boxes = yolo2xyxy(boxes[..., 0:4])
+
+        labels = boxes[..., 4:5]
+        instance.labels = labels
+        return instance