supervisely 6.73.357__py3-none-any.whl → 6.73.358__py3-none-any.whl

supervisely/nn/inference/tracking/tracker_interface.py

@@ -637,6 +637,29 @@ class TrackerInterfaceV2:
                 logger.info("Task stopped by user.", extra=self.log_extra)
                 self.stop()
 
+    def notify_error(self, exception: Exception):
+        logger.debug(f"Notify error: {str(exception)}", extra=self.log_extra)
+        error = type(exception).__name__
+        message = str(exception)
+        if self.direct_progress:
+            self.api.vid_ann_tool.set_direct_tracking_error(
+                self.session_id,
+                self.video_id,
+                self.track_id,
+                f"{type(exception).__name__}: {str(exception)}",
+            )
+        else:
+            self.api.video.notify_tracking_error(self.track_id, error, message)
+
+    def notify_warning(self, message: str):
+        logger.debug(f"Notify warning: {message}", extra=self.log_extra)
+        if self.direct_progress:
+            self.api.vid_ann_tool.set_direct_tracking_warning(
+                self.session_id, self.video_id, self.track_id, message
+            )
+        else:
+            self.api.video.notify_tracking_warning(self.track_id, self.video_id, message)
+
     def _upload_exception_handler(self, exception: Exception):
         logger.error(
             "Error in upload loop: %s", str(exception), exc_info=True, extra=self.log_extra

supervisely/nn/inference/uploader.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import threading
+from logging import Logger
+from queue import Empty, Queue
+from types import TracebackType
+from typing import Callable, Optional, Type
+
+from supervisely.sly_logger import logger
+
+
+class Uploader:
+    def __init__(
+        self,
+        upload_f: Callable,
+        notify_f: Callable = None,
+        exception_handler: Callable = None,
+        logger: Logger = None,
+    ):
+        self._upload_f = upload_f
+        self._notify_f = notify_f
+        self._exception_handler = exception_handler
+        if self._exception_handler is None:
+            self._exception_handler = self._default_exception_handler
+        self._logger = logger
+        self.exception = None
+        self._lock = threading.Lock()
+        self._q = Queue()
+        self._stop_event = threading.Event()
+        self._exception_event = threading.Event()
+        self._upload_thread = threading.Thread(
+            target=self._upload_loop,
+            daemon=True,
+        )
+        self.raise_from_notify = False
+        self._notify_thread = None
+        self._notify_q = Queue()
+        if self._notify_f is not None:
+            self._notify_thread = threading.Thread(target=self._notify_loop, daemon=True)
+            self._notify_thread.start()
+        self._upload_thread.start()
+
+    def _notify_loop(self):
+        try:
+            while (
+                not self._stop_event.is_set() or not self._notify_q.empty() or not self._q.empty()
+            ):
+                if self._exception_event.is_set():
+                    return
+                items = []
+                try:
+                    timeout = 0.1 if self._stop_event.is_set() else 1.0
+                    item = self._notify_q.get(timeout=timeout)
+                    items.append(item)
+                    while True:
+                        try:
+                            items.append(self._notify_q.get_nowait())
+                        except Empty:
+                            break
+                    if items:
+                        self._notify_f(items)
+
+                    for _ in range(len(items)):
+                        self._notify_q.task_done()
+                except Empty:
+                    continue
+        except StopIteration:
+            self.stop()
+            return
+        except Exception as e:
+            try:
+                raise RuntimeError("Error in notify loop") from e
+            except RuntimeError as e_:
+                e = e_
+            if self._logger is not None:
+                self._logger.error("Error in notify loop: %s", str(e), exc_info=True)
+            if self.raise_from_notify and not self._exception_event.is_set():
+                self.set_exception(e)
+            return
+
+    def _upload_loop(self):
+        try:
+            while not self._stop_event.is_set() or not self._q.empty():
+                if self._exception_event.is_set():
+                    return
+                items = []
+                try:
+                    timeout = 0.1 if self._stop_event.is_set() else 1.0
+                    item = self._q.get(timeout=timeout)
+                    items.append(item)
+                    while True:
+                        try:
+                            items.append(self._q.get_nowait())
+                        except Empty:
+                            break
+                    if items:
+                        self._upload_f(items)
+                        self.notify(items)
+
+                    for _ in range(len(items)):
+                        self._q.task_done()
+                except Empty:
+                    continue
+        except StopIteration:
+            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():
+                self.set_exception(e)
+            return
+
+    def put(self, items):
+        for item in items:
+            self._q.put(item)
+
+    def notify(self, items):
+        if self._notify_thread is not None:
+            for item in items:
+                self._notify_q.put(item)
+
+    def stop(self):
+        self._stop_event.set()
+
+    def join(self, timeout=None):
+        self._upload_thread.join(timeout=timeout)
+        if self._notify_thread is not None:
+            self._notify_thread.join(timeout=timeout)
+
+    def has_exception(self):
+        return self._exception_event.is_set()
+
+    def set_exception(self, exception: Exception):
+        with self._lock:
+            self._exception_event.set()
+            self.exception = exception
+
+    def __enter__(self):
+        return self
+
+    def _default_exception_handler(
+        self,
+        exception: Exception,
+    ):
+        raise exception
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+        try:
+            self.join(timeout=5)
+        except TimeoutError:
+            _logger = logger
+            if self._logger is not None:
+                _logger = self._logger
+            _logger.warning("Uploader thread didn't finish in time")
+        if exc_type is not None:
+            exc = exc_val.with_traceback(exc_tb)
+            return self._exception_handler(exc)
+        return False

supervisely/nn/model/model_api.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+import os
+from os import PathLike
+from typing import List, Union
+
+import numpy as np
+import requests
+
+import supervisely.io.env as sly_env
+import supervisely.io.json as sly_json
+from supervisely.api.module_api import ApiField
+from supervisely.api.task_api import TaskApi
+from supervisely.nn.experiments import ExperimentInfo
+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:
+    def __init__(self, api: "Api" = None, task_id: int = None, url: str = None):
+        assert not (task_id is None and url is None), "Either `task_id` or `url` must be passed."
+        assert (
+            task_id is None or url is None
+        ), "Either `task_id` or `url` must be passed (not both)."
+        if task_id is not None:
+            assert api is not None, "API must be provided if `task_id` is passed."
+
+        self.api = api
+        self.task_id = task_id
+        self.url = url
+
+        if self.task_id is not None:
+            task_info = self.api.task.get_info_by_id(self.task_id)
+            if task_info is None:
+                raise ValueError(f"Task with id {self.task_id} not found.")
+            self.url = f'{self.api.server_address}/net/{task_info["meta"]["sessionToken"]}'
+
+    # region Requests
+    def _post(self, method: str, data: dict, raise_for_status: bool = True):
+        url = f"{self.url.rstrip('/')}/{method.lstrip('/')}"
+        response = requests.post(url, json=data)
+        if raise_for_status:
+            response.raise_for_status()
+        return response.json()
+
+    def _get(self, method: str, params: dict = None, raise_for_status: bool = True):
+        url = f"{self.url.rstrip('/')}/{method.lstrip('/')}"
+        response = requests.get(url, params=params)
+        if raise_for_status:
+            response.raise_for_status()
+        return response.json()
+
+    # ------------------------------------ #
+
+    # region Info
+    def get_info(self):
+        if self.task_id is not None:
+            return self.api.nn._deploy_api.get_deploy_info(self.task_id)
+        return self._post("get_deploy_info", {})
+
+    def get_settings(self):
+        if self.task_id is not None:
+            return self.api.task.send_request(self.task_id, "get_custom_inference_settings", {})[
+                "settings"
+            ]
+        else:
+            return self._post("get_custom_inference_settings", {})["settings"]
+
+    def get_model_meta(self):
+        if self.task_id is not None:
+            return ProjectMeta.from_json(
+                self.api.task.send_request(self.task_id, "get_output_classes_and_tags", {})
+            )
+        else:
+            return ProjectMeta.from_json(self._post("get_output_classes_and_tags", {}))
+
+    def get_classes(self):
+        model_meta = self.get_model_meta()
+        return [obj_class.name for obj_class in model_meta.obj_classes]
+
+    def list_pretrained_models(self) -> List[str]:
+        """Return a list of pretrained model names available for deployment"""
+        return self._post("list_pretrained_models", {})
+
+    def list_pretrained_model_infos(self) -> List[dict]:
+        """Return a list of pretrained model infos with full information about each model"""
+        return self._post("list_pretrained_model_infos", {})
+
+    def list_experiments(self) -> List[ExperimentInfo]:
+        """Return a list of training experiments in Supervisely"""
+        raise NotImplementedError
+
+    def is_deployed(self) -> bool:
+        if self.task_id is not None:
+            return self.api.task.is_ready(self.task_id)
+        return self._post("is_ready", {})["status"] == "ready"
+
+    def status(self):
+        if self.task_id is not None:
+            return self.api.task.send_request(self.task_id, "get_status", {})
+        return self._post("get_status", {})
+
+    def shutdown(self):
+        if self.task_id is not None:
+            return self.api.task.stop(self.task_id)
+        response = self._post("tasks.stop", {ApiField.ID: id})
+        return TaskApi.Status(response[ApiField.STATUS])
+
+    # --------------------- #
+
+    # region Load
+    def load(
+        self,
+        model: str,
+        device: str = None,
+        runtime: str = None,
+    ):
+        if self.url is not None:
+            if os.path.exists(model):
+                self._load_local_custom_model(model, device, runtime)
+            else:
+                self._load_local_pretrained_model(model, device, runtime)
+
+        elif model.startswith("/"):
+            team_id = sly_env.team_id()
+            artifacts_dir, checkpoint_name = (
+                self.api.nn._deploy_api._get_artifacts_dir_and_checkpoint_name(model)
+            )
+            self.api.nn._deploy_api.load_custom_model(
+                self.task_id,
+                team_id,
+                artifacts_dir,
+                checkpoint_name,
+                device=device,
+                runtime=runtime,
+            )
+        else:
+            self.api.nn._deploy_api.load_pretrained_model(
+                self.task_id, model, device=device, runtime=runtime
+            )
+
+    def _load_local_pretrained_model(self, model: str, device: str = None, runtime: str = None):
+        available_models = self.list_pretrained_models()
+        if model not in available_models:
+            raise ValueError(f"Model {model} not found in available models: {available_models}")
+
+        deploy_params = {
+            "model_files": {},
+            "model_source": ModelSource.PRETRAINED,
+            "model_info": {},
+            "device": device,
+            "runtime": runtime,
+        }
+        state = {"deploy_params": deploy_params, "model_name": model}
+        return self._post("deploy_from_api", {"state": state})
+
+    def _load_local_custom_model(self, model: str, device: str = None, runtime: str = None):
+        deploy_params = self._get_custom_model_params(model, device, runtime)
+        state = {"deploy_params": deploy_params, "model_name": model}
+        return self._post("deploy_from_api", {"state": state})
+
+    def _get_custom_model_params(self, model_name: str, device: str = None, runtime: str = None):
+        def _load_experiment_info(artifacts_dir):
+            experiment_path = os.path.join(artifacts_dir, "experiment_info.json")
+            model_info = sly_json.load_json_file(experiment_path)
+            model_meta_path = os.path.join(artifacts_dir, "model_meta.json")
+            model_info["model_meta"] = sly_json.load_json_file(model_meta_path)
+            original_model_files = model_info.get("model_files")
+            return model_info, original_model_files
+
+        def _prepare_local_model_files(artifacts_dir, checkpoint_path, original_model_files):
+            return {k: os.path.join(artifacts_dir, v) for k, v in original_model_files.items()} | {
+                "checkpoint": checkpoint_path
+            }
+
+        model_source = ModelSource.CUSTOM
+        artifacts_dir = os.path.dirname(os.path.dirname(model_name))
+        model_info, original_model_files = _load_experiment_info(artifacts_dir)
+        model_files = _prepare_local_model_files(artifacts_dir, model_name, original_model_files)
+        deploy_params = {
+            "model_files": model_files,
+            "model_source": model_source,
+            "model_info": model_info,
+            "device": device,
+            "runtime": runtime,
+        }
+        return deploy_params
+
+    # --------------------------------- #
+
+    # region Predict
+    def predict_detached(
+        self,
+        input: Union[np.ndarray, str, PathLike, list] = None,
+        image_id: int = None,
+        video_id: int = None,
+        dataset_id: int = None,
+        project_id: int = None,
+        batch_size: int = None,
+        conf: float = None,
+        classes: List[str] = None,
+        upload_mode: str = None,
+        **kwargs,
+    ) -> PredictionSession:
+        if upload_mode is not None:
+            kwargs["upload_mode"] = upload_mode
+        return 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,
+            classes=classes,
+            **kwargs,
+        )
+
+    def predict(
+        self,
+        input: Union[np.ndarray, str, PathLike, list] = None,
+        image_id: Union[List[int], int] = None,
+        video_id: Union[List[int], int] = None,
+        dataset_id: Union[List[int], int] = None,
+        project_id: Union[List[int], int] = None,
+        batch_size: int = None,
+        conf: float = None,
+        img_size: int = None,
+        classes: List[str] = None,
+        upload_mode: str = None,
+        recursive: bool = 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,
+            )
+        )
+
+    # ------------------------------------ #