supervisely 6.73.417__py3-none-any.whl → 6.73.419__py3-none-any.whl

supervisely/api/entity_annotation/figure_api.py

@@ -800,6 +800,7 @@ class FigureApi(RemoveableBulkModuleApi):
         skip_geometry: bool = False,
         semaphore: Optional[asyncio.Semaphore] = None,
         log_progress: bool = True,
+        batch_size: int = 300,
     ) -> Dict[int, List[FigureInfo]]:
         """
         Asynchronously download figures for the given dataset ID. Can be filtered by image IDs.
@@ -815,6 +816,10 @@ class FigureApi(RemoveableBulkModuleApi):
         :type semaphore: Optional[asyncio.Semaphore], optional
         :param log_progress: If True, log the progress of the download.
         :type log_progress: bool, optional
+        :param batch_size: Size of the batch for downloading figures per 1 request. Default is 300.
+                        Used for batching image_ids when filtering by specific images.
+                        Adjust this value for optimal performance, value cannot exceed 500.
+        :type batch_size: int, optional
         :return: A dictionary where keys are image IDs and values are lists of figures.
         :rtype: Dict[int, List[FigureInfo]]
 
@@ -853,71 +858,104 @@ class FigureApi(RemoveableBulkModuleApi):
         if skip_geometry is True:
             fields = [x for x in fields if x != ApiField.GEOMETRY]
 
-        if image_ids is None:
-            filters = []
-        else:
-            filters = [
-                {
-                    ApiField.FIELD: ApiField.ENTITY_ID,
-                    ApiField.OPERATOR: "in",
-                    ApiField.VALUE: image_ids,
-                }
-            ]
-
-        data = {
+        # Base data setup
+        base_data = {
             ApiField.DATASET_ID: dataset_id,
             ApiField.FIELDS: fields,
-            ApiField.FILTER: filters,
         }
 
-        # Get first page to determine total pages
         if semaphore is None:
             semaphore = self._api.get_default_semaphore()
-        images_figures = defaultdict(list)
-        pages_count = None
-        total = 0
-        tasks = []
 
-        async def _get_page(page_data, page_num):
+        async def _get_page_figures(page_data, semaphore, progress_cb: tqdm = None):
+            """Helper function to get figures from a single page"""
             async with semaphore:
                 response = await self._api.post_async("figures.list", page_data)
                 response_json = response.json()
-                nonlocal pages_count, total
-                pages_count = response_json["pagesCount"]
-                if page_num == 1:
-                    total = response_json["total"]
 
                 page_figures = []
                 for info in response_json["entities"]:
                     figure_info = self._convert_json_info(info, True)
                     page_figures.append(figure_info)
+                if progress_cb is not None:
+                    progress_cb.update(len(response_json["entities"]))
                 return page_figures
 
-        # Get first page
-        data[ApiField.PAGE] = 1
-        first_page_figures = await _get_page(data, 1)
+        async def _get_all_pages(ids_filter, progress_cb: tqdm = None):
+            """Internal function to process all pages for given filter"""
+            data = base_data.copy()
+            data[ApiField.FILTER] = ids_filter
+
+            # Get first page to determine pagination
+            data[ApiField.PAGE] = 1
+            async with semaphore:
+                response = await self._api.post_async("figures.list", data)
+                response_json = response.json()
+
+            pages_count = response_json["pagesCount"]
+            all_figures = []
+
+            # Process first page
+            for info in response_json["entities"]:
+                figure_info = self._convert_json_info(info, True)
+                all_figures.append(figure_info)
+            if progress_cb is not None:
+                progress_cb.update(len(response_json["entities"]))
+
+            # Process remaining pages in parallel if needed
+            if pages_count > 1:
+                tasks = []
+                for page in range(2, pages_count + 1):
+                    page_data = data.copy()
+                    page_data[ApiField.PAGE] = page
+                    tasks.append(
+                        asyncio.create_task(
+                            _get_page_figures(page_data, semaphore, progress_cb=progress_cb)
+                        )
+                    )
+
+                if tasks:
+                    page_results = await asyncio.gather(*tasks)
+                    for page_figures in page_results:
+                        all_figures.extend(page_figures)
+
+            return all_figures
 
         if log_progress:
-            progress_cb = tqdm(total=total, desc="Downloading figures")
+            progress_cb = tqdm(desc="Downloading figures", unit="figure", total=0)
+        else:
+            progress_cb = None
 
-        for figure in first_page_figures:
-            images_figures[figure.entity_id].append(figure)
-            if log_progress:
-                progress_cb.update(1)
-
-        # Get rest of the pages in parallel
-        if pages_count > 1:
-            for page in range(2, pages_count + 1):
-                page_data = data.copy()
-                page_data[ApiField.PAGE] = page
-                tasks.append(asyncio.create_task(_get_page(page_data, page)))
-
-            for task in asyncio.as_completed(tasks):
-                page_figures = await task
-                for figure in page_figures:
-                    images_figures[figure.entity_id].append(figure)
-                    if log_progress:
-                        progress_cb.update(1)
+        # Strategy: batch processing based on image_ids
+        tasks = []
+
+        if image_ids is None:
+            # Single task for all figures in dataset
+            filters = []
+            tasks.append(_get_all_pages(filters, progress_cb=progress_cb))
+        else:
+            # Batch image_ids and create tasks for each batch
+            for batch_ids in batched(image_ids, batch_size):
+                filters = [
+                    {
+                        ApiField.FIELD: ApiField.ENTITY_ID,
+                        ApiField.OPERATOR: "in",
+                        ApiField.VALUE: list(batch_ids),
+                    }
+                ]
+                tasks.append(_get_all_pages(filters, progress_cb=progress_cb))
+                # Small delay between batches to reduce server load
+                await asyncio.sleep(0.02)
+
+        # Execute all tasks in parallel and collect results
+        all_results = await asyncio.gather(*tasks)
+
+        # Combine results from all batches
+        images_figures = defaultdict(list)
+
+        for batch_figures in all_results:
+            for figure in batch_figures:
+                images_figures[figure.entity_id].append(figure)
 
         return dict(images_figures)
 
@@ -928,6 +966,7 @@ class FigureApi(RemoveableBulkModuleApi):
         skip_geometry: bool = False,
         semaphore: Optional[asyncio.Semaphore] = None,
         log_progress: bool = True,
+        batch_size: int = 300,
     ) -> Dict[int, List[FigureInfo]]:
         """
         Download figures for the given dataset ID. Can be filtered by image IDs.
@@ -945,6 +984,10 @@ class FigureApi(RemoveableBulkModuleApi):
         :type semaphore: Optional[asyncio.Semaphore], optional
         :param log_progress: If True, log the progress of the download.
         :type log_progress: bool, optional
+        :param batch_size: Size of the batch for downloading figures per 1 request. Default is 300.
+                        Used for batching image_ids when filtering by specific images.
+                        Adjust this value for optimal performance, value cannot exceed 500.
+        :type batch_size: int, optional
 
         :return: A dictionary where keys are image IDs and values are lists of figures.
         :rtype: Dict[int, List[FigureInfo]]
@@ -970,6 +1013,7 @@ class FigureApi(RemoveableBulkModuleApi):
                     skip_geometry=skip_geometry,
                     semaphore=semaphore,
                     log_progress=log_progress,
+                    batch_size=batch_size,
                 )
             )
         except Exception:

supervisely/nn/inference/inference.py

@@ -1265,6 +1265,26 @@ class Inference:
 
     def get_classes(self) -> List[str]:
         return self.classes
+    
+    def _tracker_init(self, tracker: str, tracker_settings: dict):
+        # Check if tracking is supported for this model
+        info = self.get_info()
+        tracking_support = info.get("tracking_on_videos_support", False)
+        
+        if not tracking_support:
+            logger.debug("Tracking is not supported for this model")
+            return None
+        
+        if tracker == "botsort":
+            from supervisely.nn.tracker import BotSortTracker
+            device = tracker_settings.get("device", self.device)
+            logger.debug(f"Initializing BotSort tracker with device: {device}")
+            return BotSortTracker(settings=tracker_settings, device=device)
+        else:
+            if tracker is not None:
+                logger.warning(f"Unknown tracking type: {tracker}. Tracking is disabled.")
+            return None
+
 
     def get_info(self) -> Dict[str, Any]:
         num_classes = None
@@ -1291,9 +1311,9 @@ class Inference:
             "sliding_window_support": self.sliding_window_mode,
             "videos_support": True,
             "async_video_inference_support": True,
-            "tracking_on_videos_support": True,
+            "tracking_on_videos_support": False,
             "async_image_inference_support": True,
-            "tracking_algorithms": ["bot", "deepsort"],
+            "tracking_algorithms": ["botsort"],
             "batch_inference_support": self.is_batch_inference_supported(),
             "max_batch_size": self.max_batch_size,
         }
@@ -1847,24 +1867,12 @@ class Inference:
         else:
             n_frames = frames_reader.frames_count()
 
-        if tracking == "bot":
-            from supervisely.nn.tracker import BoTTracker
-
-            tracker = BoTTracker(state)
-        elif tracking == "deepsort":
-            from supervisely.nn.tracker import DeepSortTracker
-
-            tracker = DeepSortTracker(state)
-        else:
-            if tracking is not None:
-                logger.warning(f"Unknown tracking type: {tracking}. Tracking is disabled.")
-            tracker = None
-
+        self._tracker = self._tracker_init(state.get("tracker", None), state.get("tracker_settings", {}))
+        
         progress_total = (n_frames + step - 1) // step
         inference_request.set_stage(InferenceRequest.Stage.INFERENCE, 0, progress_total)
 
         results = []
-        tracks_data = {}
         for batch in batched(
             range(start_frame_index, start_frame_index + direction * n_frames, direction * step),
             batch_size,
@@ -1884,28 +1892,32 @@ class Inference:
                 source=frames,
                 settings=inference_settings,
             )
+            
+            if self._tracker is not None:
+                anns = self._apply_tracker_to_anns(frames, anns)
+                
             predictions = [
                 Prediction(ann, model_meta=self.model_meta, frame_index=frame_index)
                 for ann, frame_index in zip(anns, batch)
             ]
+            
             for pred, this_slides_data in zip(predictions, slides_data):
                 pred.extra_data["slides_data"] = this_slides_data
             batch_results = self._format_output(predictions)
-            if tracker is not None:
-                for frame_index, frame, ann in zip(batch, frames, anns):
-                    tracks_data = tracker.update(frame, ann, frame_index, tracks_data)
+            
             inference_request.add_results(batch_results)
             inference_request.done(len(batch_results))
             logger.debug(f"Frames {batch[0]}-{batch[-1]} done.")
         video_ann_json = None
-        if tracker is not None:
+        if self._tracker is not None:
             inference_request.set_stage("Postprocess...", 0, 1)
-            video_ann_json = tracker.get_annotation(
-                tracks_data, (video_height, video_witdth), n_frames
-            ).to_json()
+            
+            video_ann_json = self._tracker.video_annotation.to_json()
             inference_request.done()
         result = {"ann": results, "video_ann": video_ann_json}
         inference_request.final_result = result.copy()
+        return video_ann_json
+        
 
     def _inference_image_ids(
         self,
@@ -2083,18 +2095,8 @@ class Inference:
         else:
             n_frames = video_info.frames_count
 
-        if tracking == "bot":
-            from supervisely.nn.tracker import BoTTracker
-
-            tracker = BoTTracker(state)
-        elif tracking == "deepsort":
-            from supervisely.nn.tracker import DeepSortTracker
-
-            tracker = DeepSortTracker(state)
-        else:
-            if tracking is not None:
-                logger.warning(f"Unknown tracking type: {tracking}. Tracking is disabled.")
-            tracker = None
+        self._tracker = self._tracker_init(state.get("tracker", None), state.get("tracker_settings", {}))
+        
         logger.debug(
             f"Video info:",
             extra=dict(
@@ -2111,7 +2113,6 @@ class Inference:
         progress_total = (n_frames + step - 1) // step
         inference_request.set_stage(InferenceRequest.Stage.INFERENCE, 0, progress_total)
 
-        tracks_data = {}
         for batch in batched(
             range(start_frame_index, start_frame_index + direction * n_frames, direction * step),
             batch_size,
@@ -2130,6 +2131,10 @@ class Inference:
                 source=frames,
                 settings=inference_settings,
             )
+            
+            if self._tracker is not None:
+                anns = self._apply_tracker_to_anns(frames, anns)
+                
             predictions = [
                 Prediction(
                     ann,
@@ -2137,27 +2142,24 @@ class Inference:
                     frame_index=frame_index,
                     video_id=video_info.id,
                     dataset_id=video_info.dataset_id,
-                    project_id=video_info.project_id,
-                )
+                        project_id=video_info.project_id,
+                    )
                 for ann, frame_index in zip(anns, batch)
             ]
             for pred, this_slides_data in zip(predictions, slides_data):
                 pred.extra_data["slides_data"] = this_slides_data
             batch_results = self._format_output(predictions)
-            if tracker is not None:
-                for frame_index, frame, ann in zip(batch, frames, anns):
-                    tracks_data = tracker.update(frame, ann, frame_index, tracks_data)
+                    
             inference_request.add_results(batch_results)
             inference_request.done(len(batch_results))
             logger.debug(f"Frames {batch[0]}-{batch[-1]} done.")
         video_ann_json = None
-        if tracker is not None:
+        if self._tracker is not None:
             inference_request.set_stage("Postprocess...", 0, 1)
-            video_ann_json = tracker.get_annotation(
-                tracks_data, (video_info.frame_height, video_info.frame_width), n_frames
-            ).to_json()
+            video_ann_json = self._tracker.video_annotation.to_json()
             inference_request.done()
         inference_request.final_result = {"video_ann": video_ann_json}
+        return video_ann_json
 
     def _inference_project_id(self, api: Api, state: dict, inference_request: InferenceRequest):
         """Inference project images.
@@ -4117,6 +4119,20 @@ class Inference:
                 self._args.draw,
             )
 
+    def _apply_tracker_to_anns(self, frames: List[np.ndarray], anns: List[Annotation]):
+        updated_anns = []
+        for frame, ann in zip(frames, anns):
+            matches = self._tracker.update(frame, ann)
+            track_ids = [match["track_id"] for match in matches]
+            tracked_labels = [match["label"] for match in matches]
+            
+            filtered_annotation = ann.clone(
+                labels=tracked_labels,
+                custom_data=track_ids
+            )
+            updated_anns.append(filtered_annotation)
+        return updated_anns
+                
     def _add_workflow_input(self, model_source: str, model_files: dict, model_info: dict):
         if model_source == ModelSource.PRETRAINED:
             checkpoint_url = model_info["meta"]["model_files"]["checkpoint"]

supervisely/nn/inference/instance_segmentation/instance_segmentation.py

@@ -12,6 +12,7 @@ class InstanceSegmentation(Inference):
     def get_info(self) -> dict:
         info = super().get_info()
         info["task type"] = "instance segmentation"
+        info["tracking_on_videos_support"] = True
         # recommended parameters:
         # info["model_name"] = ""
         # info["checkpoint_name"] = ""

supervisely/nn/inference/object_detection/object_detection.py

@@ -17,6 +17,7 @@ class ObjectDetection(Inference):
     def get_info(self) -> dict:
         info = super().get_info()
         info["task type"] = "object detection"
+        info["tracking_on_videos_support"] = True
         # recommended parameters:
         # info["model_name"] = ""
         # info["checkpoint_name"] = ""

supervisely/nn/inference/session.py

@@ -271,7 +271,7 @@ class SessionJSON:
         start_frame_index: int = None,
         frames_count: int = None,
         frames_direction: Literal["forward", "backward"] = None,
-        tracker: Literal["bot", "deepsort"] = None,
+        tracker: Literal["botsort"] = None,
         batch_size: int = None,
     ) -> Dict[str, Any]:
         endpoint = "inference_video_id"
@@ -295,7 +295,7 @@ class SessionJSON:
         frames_direction: Literal["forward", "backward"] = None,
         process_fn=None,
         preparing_cb=None,
-        tracker: Literal["bot", "deepsort"] = None,
+        tracker: Literal["botsort"] = None,
         batch_size: int = None,
     ) -> Iterator:
         if self._async_inference_uuid:
@@ -795,7 +795,7 @@ class Session(SessionJSON):
         start_frame_index: int = None,
         frames_count: int = None,
         frames_direction: Literal["forward", "backward"] = None,
-        tracker: Literal["bot", "deepsort"] = None,
+        tracker: Literal["botsort"] = None,
         batch_size: int = None,
     ) -> List[sly.Annotation]:
         pred_list_raw = super().inference_video_id(
@@ -811,7 +811,7 @@ class Session(SessionJSON):
         start_frame_index: int = None,
         frames_count: int = None,
         frames_direction: Literal["forward", "backward"] = None,
-        tracker: Literal["bot", "deepsort"] = None,
+        tracker: Literal["botsort"] = None,
         batch_size: int = None,
         preparing_cb=None,
     ) -> AsyncInferenceIterator:

supervisely/nn/model/model_api.py

@@ -211,12 +211,15 @@ class ModelAPI:
         project_id: int = None,
         batch_size: int = None,
         conf: float = None,
+        img_size: int = None,
         classes: List[str] = None,
         upload_mode: str = None,
+        recursive: bool = False,
+        tracking: bool = None,
+        tracking_config: dict = None,
         **kwargs,
     ) -> PredictionSession:
-        if upload_mode is not None:
-            kwargs["upload_mode"] = upload_mode
+        
         return PredictionSession(
             self.url,
             input=input,
@@ -227,7 +230,12 @@ class ModelAPI:
             api=self.api,
             batch_size=batch_size,
             conf=conf,
+            img_size=img_size,
             classes=classes,
+            upload_mode=upload_mode,
+            recursive=recursive,
+            tracking=tracking,
+            tracking_config=tracking_config,
             **kwargs,
         )
 
@@ -243,28 +251,31 @@ class ModelAPI:
         img_size: int = None,
         classes: List[str] = None,
         upload_mode: str = None,
-        recursive: bool = None,
+        recursive: bool = False,
+        tracking: bool = None,
+        tracking_config: dict = None,
         **kwargs,
     ) -> List[Prediction]:
         if "show_progress" not in kwargs:
             kwargs["show_progress"] = True
-        if recursive is not None:
-            kwargs["recursive"] = recursive
-        if img_size is not None:
-            kwargs["img_size"] = img_size
-        return list(
-            self.predict_detached(
-                input,
-                image_id,
-                video_id,
-                dataset_id,
-                project_id,
-                batch_size,
-                conf,
-                classes,
-                upload_mode,
-                **kwargs,
-            )
+        session = PredictionSession(
+            self.url,
+            input=input,
+            image_id=image_id,
+            video_id=video_id,
+            dataset_id=dataset_id,
+            project_id=project_id,
+            api=self.api,
+            batch_size=batch_size,
+            conf=conf,
+            img_size=img_size,
+            classes=classes,
+            upload_mode=upload_mode,
+            recursive=recursive,
+            tracking=tracking,
+            tracking_config=tracking_config,
+            **kwargs,
         )
+        return list(session)
 
     # ------------------------------------ #

supervisely/nn/model/prediction.py

@@ -82,6 +82,7 @@ class Prediction:
         self._masks = None
         self._classes = None
         self._scores = None
+        self._track_ids = None
 
         if self.path is None and isinstance(self.source, (str, PathLike)):
             self.path = str(self.source)
@@ -125,6 +126,10 @@ class Prediction:
             )
         self._boxes = np.array(self._boxes)
         self._masks = np.array(self._masks)
+        
+        custom_data = self.annotation.custom_data
+        if custom_data and isinstance(custom_data, list) and len(custom_data) == len(self.annotation.labels):
+            self._track_ids = np.array(custom_data)
 
     @property
     def boxes(self):
@@ -178,6 +183,12 @@ class Prediction:
             obj_class.name: i for i, obj_class in enumerate(self.model_meta.obj_classes)
         }
         return np.array([cls_name_to_idx[class_name] for class_name in self.classes])
+    @property
+    def track_ids(self):
+        """Get track IDs for each detection. Returns None for detections without tracking."""
+        if self._track_ids is None:
+            self._init_geometries()
+        return self._track_ids
 
     @classmethod
     def from_json(cls, json_data: Dict, **kwargs) -> "Prediction":

supervisely/nn/model/prediction_session.py

@@ -67,8 +67,11 @@ class PredictionSession:
         dataset_id: Union[List[int], int] = None,
         project_id: Union[List[int], int] = None,
         api: "Api" = None,
+        tracking: bool = None,
+        tracking_config: dict = None,
         **kwargs: dict,
-    ):
+    ): 
+                  
         extra_input_args = ["image_ids", "video_ids", "dataset_ids", "project_ids"]
         assert (
             sum(
@@ -87,6 +90,7 @@ class PredictionSession:
             == 1
         ), "Exactly one of input, image_ids, video_id, dataset_id, project_id or image_id must be provided."
 
+        
         self._iterator = None
         self._base_url = url
         self.inference_request_uuid = None
@@ -111,6 +115,22 @@ class PredictionSession:
         self.inference_settings = {
             k: v for k, v in kwargs.items() if isinstance(v, (str, int, float))
         }
+        
+        if tracking is True:
+            model_info = self._get_session_info()
+            if not model_info.get("tracking_on_videos_support", False):
+                raise ValueError("Tracking is not supported by this model")
+            
+            if tracking_config is None:
+                self.tracker = "botsort"
+                self.tracker_settings = {}
+            else:
+                cfg = dict(tracking_config)
+                self.tracker = cfg.pop("tracker", "botsort")
+                self.tracker_settings = cfg
+        else:
+            self.tracker = None
+            self.tracker_settings = None
 
         # extra input args
         image_ids = self._set_var_from_kwargs("image_ids", kwargs, image_id)
@@ -180,7 +200,7 @@ class PredictionSession:
                         self._iterator = self._predict_images(input, **kwargs)
                     elif ext.lower() in ALLOWED_VIDEO_EXTENSIONS:
                         kwargs = get_valid_kwargs(kwargs, self._predict_videos, exclude=["videos"])
-                        self._iterator = self._predict_videos(input, **kwargs)
+                        self._iterator = self._predict_videos(input, tracker=self.tracker, tracker_settings=self.tracker_settings, **kwargs)
                     else:
                         raise ValueError(
                             f"Unsupported file extension: {ext}. Supported extensions are: {SUPPORTED_IMG_EXTS + ALLOWED_VIDEO_EXTENSIONS}"
@@ -193,7 +213,7 @@ class PredictionSession:
             if len(video_ids) > 1:
                 raise ValueError("Only one video id can be provided.")
             kwargs = get_valid_kwargs(kwargs, self._predict_videos, exclude=["videos"])
-            self._iterator = self._predict_videos(video_ids, **kwargs)
+            self._iterator = self._predict_videos(video_ids, tracker=self.tracker, tracker_settings=self.tracker_settings, **kwargs)
         elif dataset_ids is not None:
             kwargs = get_valid_kwargs(
                 kwargs,
@@ -259,7 +279,7 @@ class PredictionSession:
         if self.api is not None:
             return self.api.token
         return env.api_token(raise_not_found=False)
-
+    
     def _get_json_body(self):
         body = {"state": {}, "context": {}}
         if self.inference_request_uuid is not None:
@@ -269,7 +289,7 @@ class PredictionSession:
         if self.api_token is not None:
             body["api_token"] = self.api_token
         return body
-
+    
     def _post(self, method, *args, retries=5, **kwargs) -> requests.Response:
         if kwargs.get("headers") is None:
             kwargs["headers"] = {}
@@ -303,6 +323,11 @@ class PredictionSession:
                 if retry_idx + 1 == retries:
                     raise exc
 
+    def _get_session_info(self) -> Dict[str, Any]:
+        method = "get_session_info"
+        r = self._post(method, json=self._get_json_body())
+        return r.json()
+    
     def _get_inference_progress(self):
         method = "get_inference_progress"
         r = self._post(method, json=self._get_json_body())
@@ -558,7 +583,8 @@ class PredictionSession:
         end_frame=None,
         duration=None,
         direction: Literal["forward", "backward"] = None,
-        tracker: Literal["bot", "deepsort"] = None,
+        tracker: Literal["botsort"] = None,
+        tracker_settings: dict = None,
         batch_size: int = None,
     ):
         if len(videos) != 1:
@@ -573,6 +599,7 @@ class PredictionSession:
             ("duration", duration),
             ("direction", direction),
             ("tracker", tracker),
+            ("tracker_settings", tracker_settings), 
             ("batch_size", batch_size),
         ):
             if value is not None: