feat(client): support getting mask urls by data details

PR Closed: #1114
Graviti-AI · Dec 8, 2021 · 0561368 · 0561368
1 parent a5061cc
commit 0561368
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/tensorbay/client/segment.py b/tensorbay/client/segment.py
@@ -390,12 +390,13 @@ def _generate_data(self, offset: int = 0, limit: int = 128) -> Generator[RemoteD
 
         mask_urls = {}
         for key in _MASK_KEYS:
-            mask_urls[key] = LazyPage(
+            mask_urls[key] = LazyPage.from_items(
                 offset,
                 limit,
-                lambda offset, limit, k=key: self._generate_mask_urls(  # type: ignore[misc]
-                    k.upper(), offset, limit
+                lambda offset, limit, k=key.upper(): (  # type: ignore[misc]
+                    self._generate_mask_urls(k, offset, limit)
                 ),
+                (item["label"].get(key.upper(), {}).get("url") for item in response["dataDetails"]),
             )
 
         for i, item in enumerate(response["dataDetails"]):
@@ -408,9 +409,7 @@ def _generate_data(self, offset: int = 0, limit: int = 128) -> Generator[RemoteD
             for key in _MASK_KEYS:
                 mask = getattr(label, key, None)
                 if mask:
-                    mask.url = URL.from_getter(
-                        mask_urls[key].items[i].get, mask_urls[key].pull  # type: ignore[arg-type]
-                    )
+                    mask.url = URL.from_getter(mask_urls[key].items[i].get, mask_urls[key].pull)
                     mask.cache_path = os.path.join(self._cache_path, key, mask.path)
 
             yield data
@@ -711,8 +710,8 @@ def get_data(self, remote_path: str) -> RemoteData:
         for key in _MASK_KEYS:
             mask = getattr(label, key, None)
             if mask:
-                mask.url = URL.from_getter(
-                    lambda k=key.upper(), r=remote_path: self._get_mask_url(k, r),
+                mask.url = URL(
+                    data_details["label"][key.upper()]["url"],
                     lambda k=key.upper(), r=remote_path: (  # type: ignore[misc, arg-type]
                         self._get_mask_url(k, r)
                     ),