From 056136840b8dda6a60278f48554f40f1fea4fca7 Mon Sep 17 00:00:00 2001 From: "changjun.zhu" Date: Mon, 22 Nov 2021 11:51:06 +0800 Subject: [PATCH] feat(client): support getting mask urls by data details PR Closed: https://github.com/Graviti-AI/tensorbay-python-sdk/pull/1114 --- tensorbay/client/segment.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tensorbay/client/segment.py b/tensorbay/client/segment.py index 7f51c7564..63e526bd5 100644 --- 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) ),