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

supervisely/nn/inference/tracking/bbox_tracking.py

@@ -1,8 +1,6 @@
 import time
 import uuid
 from pathlib import Path
-from queue import Queue
-from threading import Event, Thread
 from typing import Any, BinaryIO, Dict, List, Optional
 
 import numpy as np
@@ -17,14 +15,16 @@ from supervisely.api.video.video_figure_api import FigureInfo
 from supervisely.geometry.helpers import deserialize_geometry
 from supervisely.geometry.rectangle import Rectangle
 from supervisely.imaging import image as sly_image
+from supervisely.nn.inference.inference import Uploader
+from supervisely.nn.inference.inference_request import InferenceRequest
 from supervisely.nn.inference.tracking.base_tracking import BaseTracking
 from supervisely.nn.inference.tracking.tracker_interface import (
     TrackerInterface,
     TrackerInterfaceV2,
 )
+from supervisely.nn.inference.uploader import Uploader
 from supervisely.nn.prediction_dto import Prediction, PredictionBBox
 from supervisely.sly_logger import logger
-from supervisely.task.progress import Progress
 
 
 class BBoxTracking(BaseTracking):
@@ -33,14 +33,19 @@ class BBoxTracking(BaseTracking):
         geometry_json = data["data"]
         return deserialize_geometry(geometry_type_str, geometry_json)
 
-    def _track(self, api: Api, context: dict, notify_annotation_tool: bool):
+    def _track(
+        self,
+        api: Api,
+        context: dict,
+        inference_request: InferenceRequest,
+    ):
         video_interface = TrackerInterface(
             context=context,
             api=api,
             load_all_frames=False,
             frame_loader=self.cache.download_frame,
             frames_loader=self.cache.download_frames,
-            should_notify=notify_annotation_tool,
+            should_notify=True,
         )
 
         range_of_frames = [
@@ -62,36 +67,41 @@ class BBoxTracking(BaseTracking):
                 video_id=video_interface.video_id,
             )
 
-        api.logger.info("Start tracking.")
+        def _upload_f(items: List):
+            video_interface.add_object_geometries_on_frames(*list(zip(*items)), notify=False)
+            inference_request.done(len(items))
+
+        def _notify_f(items: List):
+            frame_range = [
+                min(frame_index for (_, _, frame_index) in items),
+                max(frame_index for (_, _, frame_index) in items),
+            ]
+            pos_inc = inference_request.progress.current - video_interface.global_pos
+
+            video_interface._notify(
+                pos_increment=pos_inc,
+                fstart=frame_range[0],
+                fend=frame_range[1],
+                task=inference_request.stage,
+            )
 
-        def _upload_loop(q: Queue, stop_event: Event, video_interface: TrackerInterface):
+        def _exception_handler(exception: Exception):
             try:
-                while True:
-                    items = []
-                    while not q.empty():
-                        items.append(q.get_nowait())
-                    if len(items) > 0:
-                        video_interface.add_object_geometries_on_frames(*list(zip(*items)))
-                        continue
-                    if stop_event.is_set():
-                        video_interface._notify(True, task="stop tracking")
-                        return
-                    time.sleep(1)
-            except Exception as e:
-                api.logger.error("Error in upload loop: %s", str(e), exc_info=True)
-                video_interface._notify(True, task="stop tracking")
-                video_interface.global_stop_indicatior = True
-                raise
-
-        upload_queue = Queue()
-        stop_upload_event = Event()
-        Thread(
-            target=_upload_loop,
-            args=[upload_queue, stop_upload_event, video_interface],
-            daemon=True,
-        ).start()
-
-        try:
+                raise exception
+            except Exception:
+                api.logger.error(f"Error: {str(exception)}", exc_info=True)
+            video_interface._notify(True, task="Stop tracking due to an error")
+            raise exception
+
+        api.logger.info("Start tracking.")
+        total_progress = video_interface.frames_count * len(video_interface.figure_ids)
+        inference_request.set_stage(InferenceRequest.Stage.INFERENCE, 0, total_progress)
+        with Uploader(
+            upload_f=_upload_f,
+            notify_f=_notify_f,
+            exception_handler=_exception_handler,
+            logger=api.logger,
+        ) as uploader:
             for fig_id, obj_id in zip(
                 video_interface.geometries.keys(),
                 video_interface.object_ids,
@@ -100,7 +110,6 @@ class BBoxTracking(BaseTracking):
                 for _ in video_interface.frames_loader_generator():
                     geom = video_interface.geometries[fig_id]
                     if not isinstance(geom, Rectangle):
-                        stop_upload_event.set()
                         raise TypeError(f"Tracking does not work with {geom.geometry_name()}.")
 
                     imgs = video_interface.frames
@@ -121,21 +130,38 @@ class BBoxTracking(BaseTracking):
                         settings=self.custom_inference_settings_dict,
                     )
                     sly_geometry = self._to_sly_geometry(geometry)
-                    upload_queue.put(
-                        (sly_geometry, obj_id, video_interface._cur_frames_indexes[-1])
-                    )
 
-                    if video_interface.global_stop_indicatior:
-                        stop_upload_event.set()
-                        return
+                    uploader.put([(sly_geometry, obj_id, video_interface._cur_frames_indexes[-1])])
 
-                api.logger.info(f"Figure #{fig_id} tracked.")
-        except Exception:
-            stop_upload_event.set()
-            raise
-        stop_upload_event.set()
+                    if inference_request.is_stopped() or video_interface.global_stop_indicatior:
+                        api.logger.info(
+                            "Inference request stopped.",
+                            extra={"inference_request_uuid": inference_request.uuid},
+                        )
+                        video_interface._notify(True, task="Stop tracking")
+                        return
+                    if uploader.has_exception():
+                        exception = uploader.exception
+                        if not isinstance(exception, Exception):
+                            raise RuntimeError(
+                                f"Uploader exception is not an instance of Exception: {str(exception)}"
+                            )
+                        raise uploader.exception
+
+                api.logger.info(
+                    f"Figure #{fig_id} tracked.",
+                    extra={
+                        "figure_id": fig_id,
+                        "object_id": obj_id,
+                        "inference_request_uuid": inference_request.uuid,
+                    },
+                )
+            api.logger.info(
+                "Finished tracking.", extra={"inference_request_uuid": inference_request.uuid}
+            )
+            video_interface._notify(True, task="Finished tracking")
 
-    def _track_api(self, api: Api, context: dict, request_uuid: str = None):
+    def _track_api(self, api: Api, context: dict, inference_request: InferenceRequest):
         track_t = time.monotonic()
         # unused fields:
         context["trackId"] = "auto"
@@ -179,6 +205,8 @@ class BBoxTracking(BaseTracking):
         frames_n = video_interface.frames_count
         box_n = len(input_bboxes)
         geom_t = time.monotonic()
+
+        inference_request.set_stage(InferenceRequest.Stage.INFERENCE, 0, frames_n * box_n)
         api.logger.info(
             "Start tracking.",
             extra={
@@ -186,7 +214,7 @@ class BBoxTracking(BaseTracking):
                 "frame_range": range_of_frames,
                 "geometries_count": box_n,
                 "frames_count": frames_n,
-                "request_uuid": request_uuid,
+                "request_uuid": inference_request.uuid,
             },
         )
         for box_i, input_geom in enumerate(input_bboxes, 1):
@@ -218,6 +246,7 @@ class BBoxTracking(BaseTracking):
                 predictions_for_object.append(
                     {"type": sly_geometry.geometry_name(), "data": sly_geometry.to_json()}
                 )
+                inference_request.done()
                 api.logger.debug(
                     "Frame processed. Geometry: [%d / %d]. Frame: [%d / %d]",
                     box_i,
@@ -228,7 +257,7 @@ class BBoxTracking(BaseTracking):
                         "geometry_index": box_i,
                         "frame_index": frame_i,
                         "processing_time": time.monotonic() - frame_t,
-                        "request_uuid": request_uuid,
+                        "inference_request_uuid": inference_request.uuid,
                     },
                 )
                 frame_t = time.monotonic()
@@ -241,7 +270,7 @@ class BBoxTracking(BaseTracking):
                 extra={
                     "geometry_index": box_i,
                     "processing_time": time.monotonic() - geom_t,
-                    "request_uuid": request_uuid,
+                    "inference_request_uuid": inference_request.uuid,
                 },
             )
             geom_t = time.monotonic()
@@ -250,9 +279,12 @@ class BBoxTracking(BaseTracking):
         predictions = list(map(list, zip(*predictions)))
         api.logger.info(
             "Tracking finished.",
-            extra={"tracking_time": time.monotonic() - track_t, "request_uuid": request_uuid},
+            extra={
+                "tracking_time": time.monotonic() - track_t,
+                "inference_request_uuid": inference_request.uuid,
+            },
         )
-        return predictions
+        inference_request.final_result = predictions
 
     def _inference(self, frames: List[np.ndarray], geometries: List[Geometry], settings: dict):
         updated_settings = {
@@ -282,127 +314,161 @@ class BBoxTracking(BaseTracking):
                 )
         return results
 
-    def _track_async(self, api: Api, context: dict, inference_request_uuid: str = None):
-        inference_request = self._inference_requests[inference_request_uuid]
+    def _track_async(self, api: Api, context: dict, inference_request: InferenceRequest):
         tracker_interface = TrackerInterfaceV2(api, context, self.cache)
-        progress: Progress = inference_request["progress"]
         frames_count = tracker_interface.frames_count
         figures = tracker_interface.figures
+        frame_range = [
+            tracker_interface.frame_indexes[0],
+            tracker_interface.frame_indexes[-1],
+        ]
+        frame_range_asc = [min(frame_range), max(frame_range)]
         progress_total = frames_count * len(figures)
-        progress.total = progress_total
 
         def _upload_f(items: List[FigureInfo]):
-            with inference_request["lock"]:
-                inference_request["pending_results"].extend(items)
+            inference_request.add_results(items)
+            inference_request.done(len(items))
 
         def _notify_f(items: List[FigureInfo]):
-            items_by_object_id: Dict[int, List[FigureInfo]] = {}
-            for item in items:
-                items_by_object_id.setdefault(item.object_id, []).append(item)
-
-            for object_id, object_items in items_by_object_id.items():
-                frame_range = [
-                    min(item.frame_index for item in object_items),
-                    max(item.frame_index for item in object_items),
-                ]
-                progress.iters_done_report(len(object_items))
-                tracker_interface.notify_progress(progress.current, progress.total, frame_range)
+            frame_range = [
+                min(item.frame_index for item in items),
+                max(item.frame_index for item in items),
+            ]
+            tracker_interface.notify_progress(
+                inference_request.progress.current, inference_request.progress.total, frame_range
+            )
 
-        api.logger.info("Start tracking.")
-        try:
-            with tracker_interface(_upload_f, _notify_f):
-                for fig_i, figure in enumerate(figures, 1):
-                    figure = api.video.figure._convert_json_info(figure)
-                    if not figure.geometry_type == Rectangle.geometry_name():
-                        raise TypeError(f"Tracking does not work with {figure.geometry_type}.")
-                    api.logger.info("figure:", extra={"figure": figure._asdict()})
-                    sly_geometry: Rectangle = deserialize_geometry(
-                        figure.geometry_type, figure.geometry
+        def _exception_handler(exception: Exception):
+            api.logger.error(f"Error saving predictions: {str(exception)}", exc_info=True)
+            tracker_interface.notify_progress(
+                inference_request.progress.current,
+                inference_request.progress.current,
+                frame_range_asc,
+            )
+            tracker_interface.notify_error(exception)
+            raise Exception
+
+        api.logger.info("Start tracking.", extra={"inference_request_uuid": inference_request.uuid})
+        inference_request.set_stage(InferenceRequest.Stage.INFERENCE, 0, progress_total)
+        with Uploader(
+            upload_f=_upload_f,
+            notify_f=_notify_f,
+            exception_handler=_exception_handler,
+            logger=api.logger,
+        ) as uploader:
+            uploader.raise_from_notify
+            for fig_i, figure in enumerate(figures, 1):
+                figure = api.video.figure._convert_json_info(figure)
+                if not figure.geometry_type == Rectangle.geometry_name():
+                    raise TypeError(f"Tracking does not work with {figure.geometry_type}.")
+                api.logger.info("figure:", extra={"figure": figure._asdict()})
+                sly_geometry: Rectangle = deserialize_geometry(
+                    figure.geometry_type, figure.geometry
+                )
+                init = False
+                for frame_i, (frame, next_frame) in enumerate(
+                    tracker_interface.frames_loader_generator(), 1
+                ):
+                    target = PredictionBBox(
+                        "",  # TODO: can this be useful?
+                        [
+                            sly_geometry.top,
+                            sly_geometry.left,
+                            sly_geometry.bottom,
+                            sly_geometry.right,
+                        ],
+                        None,
                     )
-                    init = False
-                    for frame_i, (frame, next_frame) in enumerate(
-                        tracker_interface.frames_loader_generator(), 1
-                    ):
-                        target = PredictionBBox(
-                            "",  # TODO: can this be useful?
-                            [
-                                sly_geometry.top,
-                                sly_geometry.left,
-                                sly_geometry.bottom,
-                                sly_geometry.right,
-                            ],
-                            None,
-                        )
 
-                        if not init:
-                            self.initialize(frame.image, target)
-                            init = True
-
-                        logger.debug("Start prediction")
-                        t = time.time()
-                        geometry = self.predict(
-                            rgb_image=next_frame.image,
-                            prev_rgb_image=frame.image,
-                            target_bbox=target,
-                            settings=self.custom_inference_settings_dict,
-                        )
-                        logger.debug("Prediction done. Time: %f sec", time.time() - t)
-                        sly_geometry = self._to_sly_geometry(geometry)
-
-                        figure_id = uuid.uuid5(
-                            namespace=uuid.NAMESPACE_URL, name=f"{time.time()}"
-                        ).hex
-                        result_figure = api.video.figure._convert_json_info(
-                            {
-                                ApiField.ID: figure_id,
-                                ApiField.OBJECT_ID: figure.object_id,
-                                "meta": {"frame": next_frame.frame_index},
-                                ApiField.GEOMETRY_TYPE: sly_geometry.geometry_name(),
-                                ApiField.GEOMETRY: sly_geometry.to_json(),
-                                ApiField.TRACK_ID: tracker_interface.track_id,
-                            }
-                        )
+                    if not init:
+                        self.initialize(frame.image, target)
+                        init = True
 
-                        tracker_interface.add_prediction(result_figure)
+                    logger.debug("Start prediction")
+                    t = time.time()
+                    geometry = self.predict(
+                        rgb_image=next_frame.image,
+                        prev_rgb_image=frame.image,
+                        target_bbox=target,
+                        settings=self.custom_inference_settings_dict,
+                    )
+                    logger.debug("Prediction done. Time: %f sec", time.time() - t)
+                    sly_geometry = self._to_sly_geometry(geometry)
 
-                        logger.debug(
-                            "Frame [%d / %d] processed.",
-                            frame_i,
-                            tracker_interface.frames_count,
-                        )
+                    figure_id = uuid.uuid5(namespace=uuid.NAMESPACE_URL, name=f"{time.time()}").hex
+                    result_figure = api.video.figure._convert_json_info(
+                        {
+                            ApiField.ID: figure_id,
+                            ApiField.OBJECT_ID: figure.object_id,
+                            "meta": {"frame": next_frame.frame_index},
+                            ApiField.GEOMETRY_TYPE: sly_geometry.geometry_name(),
+                            ApiField.GEOMETRY: sly_geometry.to_json(),
+                            ApiField.TRACK_ID: tracker_interface.track_id,
+                        }
+                    )
 
-                        if inference_request["cancel_inference"]:
-                            return
-                        if tracker_interface.is_stopped():
-                            reason = tracker_interface.stop_reason()
-                            if isinstance(reason, Exception):
-                                raise reason
-                            return
-
-                    api.logger.info(
-                        "Figure [%d, %d] tracked.",
-                        fig_i,
-                        len(figures),
-                        extra={"figure_id": figure.id},
+                    uploader.put([result_figure])
+
+                    logger.debug(
+                        "Frame [%d / %d] processed.",
+                        frame_i,
+                        tracker_interface.frames_count,
+                        extra={
+                            "frame_index": frame_i,
+                            "figure_index": fig_i,
+                            "inference_request_uuid": inference_request.uuid,
+                        },
                     )
-        except Exception:
-            progress.message = "Error occured during tracking"
-            raise
-        else:
-            progress.message = "Ready"
-        finally:
-            progress.set(current=0, total=1, report=True)
+
+                    if inference_request.is_stopped() or tracker_interface.is_stopped():
+                        if isinstance(tracker_interface.stop_reason(), Exception):
+                            raise tracker_interface.stop_reason()
+                        api.logger.info(
+                            "Inference request stopped.",
+                            extra={"inference_request_uuid": inference_request.uuid},
+                        )
+                        tracker_interface.notify_progress(
+                            inference_request.progress.current,
+                            inference_request.progress.current,
+                            frame_range_asc,
+                        )
+                        return
+                    if uploader.has_exception():
+                        raise uploader.exception
+
+                api.logger.info(
+                    "Figure [%d, %d] tracked.",
+                    fig_i,
+                    len(figures),
+                    extra={
+                        "figure_id": figure.id,
+                        "figure_index": fig_i,
+                        "inference_request_uuid": inference_request.uuid,
+                    },
+                )
+            api.logger.info(
+                "Finished tracking", extra={"inference_request_uuid": inference_request.uuid}
+            )
+            tracker_interface.notify_progress(
+                inference_request.progress.current,
+                inference_request.progress.current,
+                frame_range_asc,
+            )
 
     def track(self, api: Api, state: Dict, context: Dict):
         fn = self.send_error_data(api, context)(self._track)
-        self.schedule_task(fn, api, context, notify_annotation_tool=True)
+        self.inference_requests_manager.schedule_task(fn, api, context)
         return {"message": "Track task started."}
 
     def track_api(self, api: Api, state: Dict, context: Dict):
-        request_uuid = uuid.uuid5(namespace=uuid.NAMESPACE_URL, name=f"{time.time()}").hex
-        result = self._track_api(api, context, request_uuid)
-        logger.info("Track-api request processed.", extra={"request_uuid": request_uuid})
-        return result
+        inference_request, future = self.inference_requests_manager.schedule_task(
+            self._track_api, api, context
+        )
+        future.result()
+        logger.info(
+            "Track-api request processed.", extra={"inference_request_uuid": inference_request.uuid}
+        )
+        return inference_request.final_result
 
     def track_api_files(self, files: List[BinaryIO], settings: Dict):
         logger.info("Start tracking with settings:", extra={"settings": settings})
@@ -425,17 +491,16 @@ class BBoxTracking(BaseTracking):
                 f"Batch size should be less than or equal to {self.max_batch_size} for this model."
             )
 
-        inference_request_uuid = uuid.uuid5(namespace=uuid.NAMESPACE_URL, name=f"{time.time()}").hex
         fn = self.send_error_data(api, context)(self._track_async)
-        self.schedule_task(fn, api, context, inference_request_uuid=inference_request_uuid)
+        inference_request, future = self.inference_requests_manager.schedule_task(fn, api, context)
 
         logger.debug(
             "Inference has scheduled from 'track_async' endpoint",
-            extra={"inference_request_uuid": inference_request_uuid},
+            extra={"inference_request_uuid": inference_request.uuid},
         )
         return {
             "message": "Inference has started.",
-            "inference_request_uuid": inference_request_uuid,
+            "inference_request_uuid": inference_request.uuid,
         }
 
     def initialize(self, init_rgb_image: np.ndarray, target_bbox: PredictionBBox) -> None: