devide into separate files

Author: Ubuntu

deeplake/integrations/mmlab/mmseg/basedataset.py

@@ -17,26 +17,25 @@
 from mmengine.utils import is_abs
 
 
-
 class BaseDataset(Dataset):
     r"""
     @brief A modified copy of OpenMMLab's BaseDataset.
 
-    This class is a direct copy of OpenMMLab's `BaseDataset`, with modifications 
-    to remove forced filesystem initialization (`force_init`) and customize the 
-    dataset length retrieval. 
+    This class is a direct copy of OpenMMLab's `BaseDataset`, with modifications
+    to remove forced filesystem initialization (`force_init`) and customize the
+    dataset length retrieval.
 
     @note
-    - We do not use the original `BaseDataset` because it enforces local filesystem 
+    - We do not use the original `BaseDataset` because it enforces local filesystem
       dataset initialization, which is incompatible with our cloud-based dataset.
-    - Instead of relying on local file scans, this version retrieves dataset size 
+    - Instead of relying on local file scans, this version retrieves dataset size
       from a cloud storage backend.
-    
+
     @modifications
     - Removed `force_init` to avoid mandatory filesystem checks.
     - Overridden `__len__` to use cloud metadata instead of local file counting.
-    
-    This ensures that the dataset can be loaded dynamically from the cloud without 
+
+    This ensures that the dataset can be loaded dynamically from the cloud without
     unnecessary local file system dependencies.
 
     The annotation format is shown as follows.
@@ -193,8 +192,7 @@ def get_data_info(self, idx: int) -> dict:
         return data_info
 
     def full_init(self):
-        """Load annotation file and set ``BaseDataset._fully_initialized`` to True.
-        """
+        """Load annotation file and set ``BaseDataset._fully_initialized`` to True."""
         if self._fully_initialized:
             return
 

deeplake/integrations/mmlab/mmseg/builder_patch.py

@@ -0,0 +1,94 @@
+import os
+import copy
+from typing import Any, Dict
+
+from deeplake.integrations.mmlab.mmseg.registry import TRANSFORMS
+from deeplake.integrations.mmlab.mmseg.load_annotations import LoadAnnotations
+from mmengine.dataset import Compose
+
+from deeplake.client.config import DEEPLAKE_AUTH_TOKEN
+
+from deeplake.util.exceptions import (
+    EmptyTokenException,
+    EmptyDeeplakePathException,
+    ConflictingDatasetParametersError,
+    MissingTensorMappingError,
+)
+
+import mmengine.registry
+
+original_build_func = mmengine.registry.DATASETS.build
+
+
+def build_transform(steps):
+    from mmengine.registry.build_functions import build_from_cfg
+
+    transforms = []
+    steps_copy = copy.deepcopy(steps)
+
+    for step in steps_copy:
+        if step["type"] == "LoadAnnotations":
+            # Create LoadAnnotations instance and add to transforms list
+            kwargs = step.copy()
+            kwargs.pop("type")
+            transform = LoadAnnotations(**kwargs)
+            transforms.append(transform)
+        elif step["type"] != "LoadImageFromFile":
+            transform = build_from_cfg(step, TRANSFORMS, None)
+            transforms.append(transform)
+
+    return Compose(transforms)
+
+
+def build_func_patch(
+    cfg: Dict,
+    *args,
+    **kwargs,
+) -> Any:
+    import deeplake as dp
+
+    creds = cfg.pop("deeplake_credentials", {})
+    token = creds.pop("token", None)
+    token = token or os.environ.get(DEEPLAKE_AUTH_TOKEN)
+    if token is None:
+        raise EmptyTokenException()
+
+    ds_path = cfg.pop("deeplake_path", None)
+    if ds_path is None or not len(ds_path):
+        raise EmptyDeeplakePathException()
+
+    deeplake_ds = dp.load(ds_path, token=token, read_only=True)[0:500:1]
+    deeplake_commit = cfg.pop("deeplake_commit", None)
+    deeplake_view_id = cfg.pop("deeplake_view_id", None)
+    deeplake_query = cfg.pop("deeplake_query", None)
+
+    if deeplake_view_id and deeplake_query:
+        raise ConflictingDatasetParametersError()
+
+    if deeplake_commit:
+        deeplake_ds.checkout(deeplake_commit)
+
+    if deeplake_view_id:
+        deeplake_ds = deeplake_ds.load_view(id=deeplake_view_id)
+
+    if deeplake_query:
+        deeplake_ds = deeplake_ds.query(deeplake_query)
+
+    ds_train_tensors = cfg.pop("deeplake_tensors", {})
+
+    if "pipeline" in cfg:
+        transform_pipeline = build_transform(cfg.get("pipeline"))
+    else:
+        transform_pipeline = None
+
+    if not ds_train_tensors and not {"img", "gt_semantic_seg"}.issubset(
+        ds_train_tensors
+    ):
+        raise MissingTensorMappingError()
+
+    cfg["lazy_init"] = False
+    res = original_build_func(cfg, *args, **kwargs)
+    res.deeplake_dataset = deeplake_ds
+    res.images_tensor = ds_train_tensors.get("img")
+    res.masks_tensor = ds_train_tensors.get("gt_semantic_seg")
+    return res, transform_pipeline

deeplake/integrations/mmlab/mmseg/load_annotations.py

@@ -0,0 +1,111 @@
+import warnings
+from mmcv.transforms import LoadAnnotations as MMCV_LoadAnnotations
+
+from deeplake.integrations.mmlab.mmseg.registry import TRANSFORMS
+
+
+@TRANSFORMS.register_module()
+class LoadAnnotations(MMCV_LoadAnnotations):
+    """Load annotations for semantic segmentation provided by dataset.
+
+    The annotation format is as the following:
+
+    .. code-block:: python
+
+        {
+            # Filename of semantic segmentation ground truth file.
+            'seg_map_path': 'a/b/c'
+        }
+
+    After this module, the annotation has been changed to the format below:
+
+    .. code-block:: python
+
+        {
+            # in str
+            'seg_fields': List
+             # In uint8 type.
+            'gt_seg_map': np.ndarray (H, W)
+        }
+
+    Required Keys:
+
+    - seg_map_path (str): Path of semantic segmentation ground truth file.
+
+    Added Keys:
+
+    - seg_fields (List)
+    - gt_seg_map (np.uint8)
+
+    Args:
+        reduce_zero_label (bool, optional): Whether reduce all label value
+            by 1. Usually used for datasets where 0 is background label.
+            Defaults to None.
+        imdecode_backend (str): The image decoding backend type. The backend
+            argument for :func:``mmcv.imfrombytes``.
+            See :fun:``mmcv.imfrombytes`` for details.
+            Defaults to 'pillow'.
+        backend_args (dict): Arguments to instantiate a file backend.
+            See https://mmengine.readthedocs.io/en/latest/api/fileio.htm
+            for details. Defaults to None.
+            Notes: mmcv>=2.0.0rc4, mmengine>=0.2.0 required.
+    """
+
+    def __init__(
+        self,
+        reduce_zero_label=None,
+        backend_args=None,
+        imdecode_backend="pillow",
+    ) -> None:
+        super().__init__(
+            with_bbox=False,
+            with_label=False,
+            with_seg=True,
+            with_keypoints=False,
+            imdecode_backend=imdecode_backend,
+            backend_args=backend_args,
+        )
+        self.reduce_zero_label = reduce_zero_label
+        if self.reduce_zero_label is not None:
+            warnings.warn(
+                "`reduce_zero_label` will be deprecated, "
+                "if you would like to ignore the zero label, please "
+                "set `reduce_zero_label=True` when dataset "
+                "initialized"
+            )
+        self.imdecode_backend = imdecode_backend
+
+    def _load_seg_map(self, results: dict) -> None:
+        """Private function to load semantic segmentation annotations.
+
+        Args:
+            results (dict): Result dict from :obj:``mmcv.BaseDataset``.
+
+        Returns:
+            dict: The dict contains loaded semantic segmentation annotations.
+        """
+
+        gt_semantic_seg = results.pop("dp_seg_map", None)
+
+        # reduce zero_label
+        if self.reduce_zero_label:
+            # avoid using underflow conversion
+            gt_semantic_seg[gt_semantic_seg == 0] = 255
+            gt_semantic_seg = gt_semantic_seg - 1
+            gt_semantic_seg[gt_semantic_seg == 254] = 255
+        # modify if custom classes
+        if results.get("label_map", None) is not None:
+            # Add deep copy to solve bug of repeatedly
+            # replace `gt_semantic_seg`, which is reported in
+            # https://github.com/open-mmlab/mmsegmentation/pull/1445/
+            gt_semantic_seg_copy = gt_semantic_seg.copy()
+            for old_id, new_id in results["label_map"].items():
+                gt_semantic_seg[gt_semantic_seg_copy == old_id] = new_id
+        results["gt_seg_map"] = gt_semantic_seg
+
+    def __repr__(self) -> str:
+        repr_str = self.__class__.__name__
+        repr_str += f"(reduce_zero_label={self.reduce_zero_label}, "
+        repr_str += f"imdecode_backend='{self.imdecode_backend}', "
+        repr_str += f"backend_args={self.backend_args})"
+        return repr_str