supervisely 6.73.410__py3-none-any.whl → 6.73.470__py3-none-any.whl

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:
@@ -441,46 +441,52 @@ class SessionJSON:
         prev_current = 0
         if preparing_cb:
             # wait for inference status
-            resp = self._get_preparing_progress()
-            awaiting_preparing_progress = 0
-            break_flag = False
-            while resp.get("status") is None:
-                time.sleep(1)
-                awaiting_preparing_progress += 1
-                if awaiting_preparing_progress > 30:
-                    break_flag = True
+            try:
                 resp = self._get_preparing_progress()
-            if break_flag:
-                logger.warning(
-                    "Unable to get preparing progress. Continue without prepaing progress status."
-                )
-            if not break_flag:
-                if resp["status"] == "download_info":
+                for i in range(30):
+                    logger.info(
+                        f"Waiting for preparing progress... {30 - i} seconds left until timeout"
+                    )
+                    resp = self._get_preparing_progress()
+                    if resp.get("status") is not None:
+                        break
+                    time.sleep(1)
+                if not resp.get("status"):
+                    raise RuntimeError("Preparing progress status is not available.")
+
+                if resp.get("status") == "download_info":
+                    logger.info("Downloading infos...")
                     progress_widget = preparing_cb(
                         message="Downloading infos", total=resp["total"], unit="it"
                     )
-                while resp["status"] == "download_info":
-                    current = resp["current"]
-                    # pylint: disable=possibly-used-before-assignment
-                    progress_widget.update(current - prev_current)
-                    prev_current = current
-                    resp = self._get_preparing_progress()
-
-                if resp["status"] == "download_project":
+                    while resp["status"] == "download_info":
+                        current = resp["current"]
+                        # pylint: disable=possibly-used-before-assignment
+                        progress_widget.update(current - prev_current)
+                        prev_current = current
+                        resp = self._get_preparing_progress()
+
+                if resp.get("status") == "download_project":
+                    logger.info("Downloading project...")
                     progress_widget = preparing_cb(message="Download project", total=resp["total"])
-                while resp["status"] == "download_project":
-                    current = resp["current"]
-                    progress_widget.update(current - prev_current)
-                    prev_current = current
-                    resp = self._get_preparing_progress()
-
-                if resp["status"] == "warmup":
+                    while resp.get("status") == "download_project":
+                        current = resp["current"]
+                        progress_widget.update(current - prev_current)
+                        prev_current = current
+                        resp = self._get_preparing_progress()
+
+                if resp.get("status") == "warmup":
+                    logger.info("Running warmup...")
                     progress_widget = preparing_cb(message="Running warmup", total=resp["total"])
-                while resp["status"] == "warmup":
-                    current = resp["current"]
-                    progress_widget.update(current - prev_current)
-                    prev_current = current
-                    resp = self._get_preparing_progress()
+                    while resp.get("status") == "warmup":
+                        current = resp["current"]
+                        progress_widget.update(current - prev_current)
+                        prev_current = current
+                        resp = self._get_preparing_progress()
+            except Exception as ex:
+                logger.warning(
+                    f"An error occurred while getting preparing progress: {ex}. Continue without preparing progress status."
+                )
 
         logger.info("Inference has started:", extra={"response": resp})
         resp, has_started = self._wait_for_async_inference_start()
@@ -537,7 +543,9 @@ class SessionJSON:
         t0 = time.time()
         while not has_started and not timeout_exceeded:
             resp = self._get_inference_progress()
-            has_started = bool(resp["result"]) or resp["progress"]["total"] != 1
+            pending_results = resp.get("pending_results", None)
+            has_results = bool(pending_results)
+            has_started = bool(resp.get("result")) or resp["progress"]["total"] != 1 or has_results
             if not has_started:
                 time.sleep(delay)
             timeout_exceeded = timeout and time.time() - t0 > timeout
@@ -795,7 +803,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 +819,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/inference/tracking/bbox_tracking.py

@@ -7,7 +7,7 @@ import numpy as np
 from pydantic import ValidationError
 
 from supervisely.annotation.annotation import Annotation
-from supervisely.annotation.label import Geometry, Label
+from supervisely.annotation.label import Geometry, Label, LabelingStatus
 from supervisely.annotation.obj_class import ObjClass
 from supervisely.api.api import Api
 from supervisely.api.module_api import ApiField
@@ -588,8 +588,12 @@ class BBoxTracking(BaseTracking):
                 # for example empty mask
                 continue
             if isinstance(label, list):
+                for lb in label:
+                    lb.status = LabelingStatus.AUTO
                 labels.extend(label)
                 continue
+
+            label.status = LabelingStatus.AUTO
             labels.append(label)
 
         # create annotation with correct image resolution

supervisely/nn/inference/tracking/point_tracking.py

@@ -10,7 +10,7 @@ from pydantic import ValidationError
 
 import supervisely.nn.inference.tracking.functional as F
 from supervisely.annotation.annotation import Annotation
-from supervisely.annotation.label import Geometry, Label
+from supervisely.annotation.label import Geometry, Label, LabelingStatus
 from supervisely.annotation.obj_class import ObjClass
 from supervisely.api.api import Api
 from supervisely.api.module_api import ApiField
@@ -610,8 +610,12 @@ class PointTracking(BaseTracking):
                 # for example empty mask
                 continue
             if isinstance(label, list):
+                for lb in label:
+                    lb.status = LabelingStatus.AUTO
                 labels.extend(label)
                 continue
+
+            label.status = LabelingStatus.AUTO
             labels.append(label)
 
         # create annotation with correct image resolution

supervisely/nn/inference/tracking/tracker_interface.py

@@ -22,6 +22,7 @@ from supervisely.geometry.rectangle import Rectangle
 from supervisely.nn.inference.cache import InferenceImageCache
 from supervisely.sly_logger import logger
 from supervisely.video_annotation.key_id_map import KeyIdMap
+from supervisely.annotation.label import LabelingStatus
 
 
 class TrackerInterface:
@@ -166,6 +167,8 @@ class TrackerInterface:
                     ApiField.GEOMETRY: geometry.to_json(),
                     ApiField.META: {ApiField.FRAME: frame_index},
                     ApiField.TRACK_ID: self.track_id,
+                    ApiField.NN_CREATED: True,
+                    ApiField.NN_UPDATED: True,
                 }
                 for geometry, frame_index in geometries_frame_indexes
             ]
@@ -195,6 +198,7 @@ class TrackerInterface:
             geometry.to_json(),
             geometry.geometry_name(),
             self.track_id,
+            status=LabelingStatus.AUTO,
         )
         self.logger.debug(f"Added {geometry.geometry_name()} to frame #{frame_ind}")
         if notify:

supervisely/nn/inference/uploader.py

@@ -105,10 +105,6 @@ class Uploader:
             self.stop()
             return
         except Exception as e:
-            try:
-                raise RuntimeError("Error in upload loop") from e
-            except RuntimeError as e_:
-                e = e_
             if self._logger is not None:
                 self._logger.error("Error in upload loop: %s", str(e), exc_info=True)
             if not self._exception_event.is_set():
@@ -152,7 +148,9 @@ class Uploader:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.stop()
         try:
-            self.join(timeout=5)
+            self.join(timeout=30)
+            if self._upload_thread.is_alive():
+                raise TimeoutError("Uploader thread didn't finish in time")
         except TimeoutError:
             _logger = logger
             if self._logger is not None:
@@ -161,4 +159,10 @@ class Uploader:
         if exc_type is not None:
             exc = exc_val.with_traceback(exc_tb)
             return self._exception_handler(exc)
+        if self.has_exception():
+            exc = self.exception
+            try:
+                raise RuntimeError(f"Error in uploader loop: {str(exc)}") from exc
+            except Exception as exc:
+                return self._exception_handler(exc)
         return False

supervisely/nn/model/model_api.py

@@ -2,6 +2,7 @@
 """load and inference models"""
 
 from __future__ import annotations
+
 import os
 from os import PathLike
 from typing import List, Union
@@ -11,6 +12,7 @@ import requests
 
 import supervisely.io.env as sly_env
 import supervisely.io.json as sly_json
+from supervisely.api.api import Api
 from supervisely.api.module_api import ApiField
 from supervisely.api.task_api import TaskApi
 from supervisely.nn.experiments import ExperimentInfo
@@ -18,7 +20,6 @@ from supervisely.nn.model.prediction import Prediction
 from supervisely.nn.model.prediction_session import PredictionSession
 from supervisely.nn.utils import ModelSource
 from supervisely.project.project_meta import ProjectMeta
-from supervisely.api.api import Api
 
 
 class ModelAPI:
@@ -71,6 +72,15 @@ class ModelAPI:
         else:
             return self._post("get_custom_inference_settings", {})["settings"]
 
+    def get_tracking_settings(self):
+        # @TODO: botsort hardcoded 
+        # Add dropdown selector for tracking algorithms later
+        if self.task_id is not None:
+            return self.api.task.send_request(self.task_id, "get_tracking_settings", {})["botsort"]
+        else:
+            return self._post("get_tracking_settings", {})["botsort"]
+
+
     def get_model_meta(self):
         if self.task_id is not None:
             return ProjectMeta.from_json(
@@ -126,7 +136,8 @@ class ModelAPI:
         device: str = None,
         runtime: str = None,
     ):
-        if self.url is not None:
+        if self.task_id is None:
+            # TODO: proper check
             if os.path.exists(model):
                 self._load_local_custom_model(model, device, runtime)
             else:
@@ -209,12 +220,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,
@@ -225,7 +239,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,
         )
 
@@ -241,28 +260,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

@@ -59,6 +59,7 @@ class Prediction:
         self.source = source
         if isinstance(annotation_json, Annotation):
             annotation_json = annotation_json.to_json()
+
         self.annotation_json = annotation_json
         self.model_meta = model_meta
         if isinstance(self.model_meta, dict):
@@ -82,6 +83,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 +127,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):
@@ -152,7 +158,7 @@ class Prediction:
 
     @property
     def annotation(self) -> Annotation:
-        if self._annotation is None:
+        if self._annotation is None and self.annotation_json is not None:
             if self.model_meta is None:
                 raise ValueError("Model meta is not provided. Cannot create annotation.")
             model_meta = get_meta_from_annotation(self.annotation_json, self.model_meta)
@@ -178,6 +184,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":
@@ -229,6 +241,8 @@ class Prediction:
             if self.image_id is not None:
                 try:
                     if api is None:
+                        # TODO: raise more clarifying error in case of failing of api init
+                        # what a user should do to fix it?
                         api = Api()
                     return api.image.download_np(self.image_id)
                 except Exception as e:

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(
@@ -112,6 +115,30 @@ class PredictionSession:
             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
+
+        if "classes" in kwargs:
+            self.inference_settings["classes"] = kwargs["classes"]
+        # TODO: remove "settings", it is the same as inference_settings
+        if "settings" in kwargs:
+            self.inference_settings.update(kwargs["settings"])
+        if "inference_settings" in kwargs:
+            self.inference_settings.update(kwargs["inference_settings"])
+
         # extra input args
         image_ids = self._set_var_from_kwargs("image_ids", kwargs, image_id)
         video_ids = self._set_var_from_kwargs("video_ids", kwargs, video_id)
@@ -139,7 +166,6 @@ class PredictionSession:
             input = [input]
         if isinstance(input[0], np.ndarray):
             # input is numpy array
-            kwargs = get_valid_kwargs(kwargs, self._predict_images, exclude=["images"])
             self._predict_images(input, **kwargs)
         elif isinstance(input[0], (str, PathLike)):
             if len(input) > 1:
@@ -180,7 +206,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 +219,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,
@@ -268,6 +294,8 @@ class PredictionSession:
             body["state"]["settings"] = self.inference_settings
         if self.api_token is not None:
             body["api_token"] = self.api_token
+        if "model_prediction_suffix" in self.kwargs:
+            body["state"]["model_prediction_suffix"] = self.kwargs["model_prediction_suffix"]
         return body
 
     def _post(self, method, *args, retries=5, **kwargs) -> requests.Response:
@@ -303,6 +331,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())
@@ -331,9 +364,21 @@ class PredictionSession:
         logger.info("Inference request will be cleared on the server")
         return r.json()
 
+    def _get_final_result(self):
+        method = "get_inference_result"
+        r = self._post(
+            method,
+            json=self._get_json_body(),
+        )
+        return r.json()
+
     def _on_infernce_end(self):
         if self.inference_request_uuid is None:
             return
+        try:
+            self.final_result = self._get_final_result()
+        except Exception as e:
+            logger.debug("Failed to get final result:", exc_info=True)
         self._clear_inference_request()
 
     @property
@@ -478,18 +523,16 @@ class PredictionSession:
                 "Inference is already running. Please stop it before starting a new one."
             )
         resp = self._post(method, **kwargs).json()
-
         self.inference_request_uuid = resp["inference_request_uuid"]
-
-        logger.info(
-            "Inference has started:",
-            extra={"inference_request_uuid": resp.get("inference_request_uuid")},
-        )
         try:
             resp, has_started = self._wait_for_inference_start(tqdm=self.tqdm)
         except:
             self.stop()
             raise
+        logger.info(
+            "Inference has started:",
+            extra={"inference_request_uuid": resp.get("inference_request_uuid")},
+        )
         frame_iterator = self.Iterator(resp["progress"]["total"], self, tqdm=self.tqdm)
         return frame_iterator
 
@@ -537,7 +580,11 @@ class PredictionSession:
         return self._predict_images_bytes(images, batch_size=batch_size)
 
     def _predict_images_ids(
-        self, images: List[int], batch_size: int = None, upload_mode: str = None
+        self,
+        images: List[int],
+        batch_size: int = None,
+        upload_mode: str = None,
+        output_project_id: int = None,
     ):
         method = "inference_batch_ids_async"
         json_body = self._get_json_body()
@@ -547,6 +594,8 @@ class PredictionSession:
             state["batch_size"] = batch_size
         if upload_mode is not None:
             state["upload_mode"] = upload_mode
+        if output_project_id is not None:
+            state["output_project_id"] = output_project_id
         return self._start_inference(method, json=json_body)
 
     def _predict_videos(
@@ -558,7 +607,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 +623,7 @@ class PredictionSession:
             ("duration", duration),
             ("direction", direction),
             ("tracker", tracker),
+            ("tracker_settings", tracker_settings), 
             ("batch_size", batch_size),
         ):
             if value is not None:
@@ -594,8 +645,11 @@ class PredictionSession:
                 encoder = MultipartEncoder(fields)
                 if self.tqdm is not None:
 
+                    bytes_read = 0
                     def _callback(monitor):
-                        self.tqdm.update(monitor.bytes_read)
+                        nonlocal bytes_read
+                        self.tqdm.update(monitor.bytes_read - bytes_read)
+                        bytes_read = monitor.bytes_read
 
                     video_size = get_file_size(video_path)
                     self._update_progress(self.tqdm, "Uploading video", 0, video_size, is_size=True)
@@ -620,6 +674,7 @@ class PredictionSession:
         upload_mode: str = None,
         iou_merge_threshold: float = None,
         cache_project_on_model: bool = None,
+        output_project_id: int = None,
     ):
         if len(project_ids) != 1:
             raise ValueError("Only one project can be processed at a time.")
@@ -637,7 +692,8 @@ class PredictionSession:
             state["iou_merge_threshold"] = iou_merge_threshold
         if cache_project_on_model is not None:
             state["cache_project_on_model"] = cache_project_on_model
-
+        if output_project_id is not None:
+            state["output_project_id"] = output_project_id
         return self._start_inference(method, json=json_body)
 
     def _predict_datasets(

supervisely/nn/prediction_dto.py

@@ -3,6 +3,7 @@ from typing import List, Optional
 import numpy as np
 
 from supervisely.geometry.cuboid_3d import Cuboid3d
+from supervisely.geometry.polyline_3d import Polyline3D
 
 
 class Prediction:
@@ -81,3 +82,9 @@ class PredictionCuboid3d(Prediction):
         super(PredictionCuboid3d, self).__init__(class_name=class_name)
         self.cuboid_3d = cuboid_3d
         self.score = score
+
+
+class PredictionPolyline3D(Prediction):
+    def __init__(self, class_name: str, polyline_3d: Polyline3D):
+        super(PredictionPolyline3D, self).__init__(class_name=class_name)
+        self.polyline_3d = polyline_3d

supervisely/nn/tracker/__init__.py

@@ -1,10 +1,8 @@
-from supervisely.sly_logger import logger
-
 try:
-    from supervisely.nn.tracker.bot_sort import BoTTracker
-    from supervisely.nn.tracker.deep_sort import DeepSortTracker
+    from supervisely.nn.tracker.botsort_tracker import BotSortTracker
+    from supervisely.nn.tracker.calculate_metrics import TrackingEvaluator, evaluate
+    TRACKING_LIBS_INSTALLED = True
 except ImportError:
-    logger.error(
-        "Failed to import tracker modules. Please try install extras with 'pip install supervisely[tracking]'"
-    )
-    raise
+    TRACKING_LIBS_INSTALLED = False
+
+from supervisely.nn.tracker.visualize import TrackingVisualizer, visualize