Merge pull request #127 from Superb-AI-Suite/hotfix/v0.19.4

fix own your storage for video sdk

spb/cli_core/commands/video_label_data.py

@@ -240,20 +240,44 @@ def _download_worker(args):
             error = {'label':str(e)}
             label_error = error
         try:
-            custom_signed_url = json.loads(label.data_url)
-            base_url = custom_signed_url['base_url']
-            query = custom_signed_url['query']
-            file_infos = custom_signed_url['file_infos']
-            for idx, file_info in enumerate(file_infos, 1):
-                file_path = f'{path}/{file_info["file_name"]}'
-                ext = file_info['file_name'].split('.')[1]
-                new_path = os.path.split(file_path)[0]
-                if not os.path.exists(new_path):
-                    os.makedirs(new_path)
-                file_url = '{}image_{:08d}.{}?{}'.format(base_url, idx, ext, query)
-                with requests_retry_session() as session:
-                    r = session.get(file_url, allow_redirects=True)
-                    open(file_path, 'wb').write(r.content)
+            data_url = json.loads(label.data_url)
+
+            if type(data_url) is list:
+                for frame_idx in range(len(data_url)):
+                    url = data_url[frame_idx]
+                    try:
+                        parts = url.split("://", 1)
+                        if len(parts) < 2:
+                            raise ValueError("Invalid URL format: missing '://' separator")
+
+                        path_parts = parts[1].split("/", 2)
+                        if len(path_parts) < 3:
+                            raise ValueError("Invalid URL format: insufficient '/' parts for path extraction")
+
+                        url_path = path_parts[2].split("?", 1)[0]
+                    except ValueError as e:
+                        raise ValueError("Invalid URL format")
+                    file_path = os.path.join(path, url_path)
+                    folder, _ = os.path.split(file_path)
+                    if folder:
+                        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                    with requests_retry_session() as session:
+                        r = session.get(url, allow_redirects=True)
+                        open(file_path, 'wb').write(r.content)
+            else:
+                base_url = data_url['base_url']
+                query = data_url['query']
+                file_infos = data_url['file_infos']
+                for idx, file_info in enumerate(file_infos, 1):
+                    file_path = os.path.join(path, file_info["file_name"])
+                    ext = file_info['file_name'].split('.')[1]
+                    folder, _ = os.path.split(file_path)
+                    if folder:
+                        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                    file_url = '{}image_{:08d}.{}?{}'.format(base_url, idx, ext, query)
+                    with requests_retry_session() as session:
+                        r = session.get(file_url, allow_redirects=True)
+                        open(file_path, 'wb').write(r.content)
         except Exception as e:
             error.update({'data':str(e)})
             data_error = error

spb/pointcloud_sdk.py

@@ -58,18 +58,20 @@ def get_data_urls(self):
 
         frame_path_prefix = "frames/frame_"
         for i, frame in enumerate(frames):
+            frame_ext = frame.get("frame_file_name", None).split(".")[-1].lower()
             index = str(i+1).zfill(8)
             frame_urls = {
-                "frame_url": f"{url}{frame_path_prefix}{index}/frame_{index}.bin?{query}",
+                "frame_url": f"{url}{frame_path_prefix}{index}/frame_{index}.{frame_ext}?{query}",
                 "frame_original_file_name": frame.get("frame_file_name", None),
                 "images": []
             }
             images = frame.get("image_infos", [])
             for k, image in enumerate(images):
+                image_ext = image.get("image_file_name", None).split(".")[-1].lower()
                 image_index = str(k+1).zfill(8)
                 frame_urls["images"].append(
                     {
-                        "image_url": f"{url}{frame_path_prefix}{index}/image_{index}_{image_index}.png?{query}",
+                        "image_url": f"{url}{frame_path_prefix}{index}/image_{index}_{image_index}.{image_ext}?{query}",
                         "image_original_file_name": image.get("image_file_name", None)
                     }
                 )

spb/sdk_config.py

@@ -1,6 +1,6 @@
 SDK_NAME='spb-cli'
 SDK_AUTHOR='Superb AI Dev Team'
-SDK_VERSION='0.19.3'
+SDK_VERSION='0.19.4'
 SDK_AUTHOR_EMAIL='support@superb-ai.com'
 SDK_DESCRIPTION='Suite Standard Library'
 SDK_LICENSE='MIT'

spb/video_sdk.py

@@ -40,9 +40,12 @@ def get_frame_url(self, idx, data_url=None):
         if data_url is None:
             data_url = json.loads(self._data.data_url)
 
-        file_ext = data_url["file_infos"][idx]["file_name"].split(".")[-1].lower()
-        file_name = f"image_{(idx+1):08}.{file_ext}"
-        return f"{data_url['base_url']}{file_name}?{data_url['query']}"
+        if type(data_url) is list:
+            return data_url[idx]
+        else:
+            file_ext = data_url["file_infos"][idx]["file_name"].split(".")[-1].lower()
+            file_name = f"image_{(idx+1):08}.{file_ext}"
+            return f"{data_url['base_url']}{file_name}?{data_url['query']}"
 
     def get_frame_urls(self):
         self._describe_data_detail()
@@ -52,8 +55,13 @@ def get_frame_urls(self):
             return None
 
         data_url = json.loads(self._data.data_url)
-        for frame_idx in range(len(data_url["file_infos"])):
-            yield self.get_frame_url(frame_idx, data_url)
+
+        if type(data_url) is list:
+            for frame_idx in range(len(data_url)):
+                yield self.get_frame_url(frame_idx, data_url)
+        else:
+            for frame_idx in range(len(data_url["file_infos"])):
+                yield self.get_frame_url(frame_idx, data_url)
 
     def get_frame(self, idx):
         return self.get_frame_url(idx)
@@ -84,13 +92,34 @@ def download_video(self, download_to=None):
             print("[INFO] Downloaded to {}".format(download_to))
 
         data_url = json.loads(self._data.data_url)
-        for frame_idx, file_info in enumerate(data_url["file_infos"]):
-            url = self.get_frame_url(frame_idx, data_url)
-            retrieve_file(
-                url=url,
-                file_path=os.path.join(download_to, file_info["file_name"])
-            )
-
+
+        if type(data_url) is list:
+            for frame_idx in range(len(data_url)):
+                url = self.get_frame_url(frame_idx, data_url)
+                try:
+                    parts = url.split("://", 1)
+                    if len(parts) < 2:
+                        raise ValueError("Invalid URL format: missing '://' separator")
+
+                    path_parts = parts[1].split("/", 2)
+                    if len(path_parts) < 3:
+                        raise ValueError("Invalid URL format: insufficient '/' parts for path extraction")
+
+                    url_path = path_parts[2].split("?", 1)[0]
+
+                except ValueError as e:
+                    raise ValueError("Invalid URL format")
+                retrieve_file(
+                    url=url,
+                    file_path=os.path.join(download_to, url_path)
+                )
+        else:
+            for frame_idx, file_info in enumerate(data_url["file_infos"]):
+                url = self.get_frame_url(frame_idx, data_url)
+                retrieve_file(
+                    url=url,
+                    file_path=os.path.join(download_to, file_info["file_name"])
+                )
         return True
 
     def get_frames(self):