supervisely 6.73.442__py3-none-any.whl → 6.73.443__py3-none-any.whl

supervisely/nn/inference/inference.py

@@ -863,6 +863,50 @@ class Inference:
                 self.gui.download_progress.hide()
         return local_model_files
 
+    def _fallback_download_custom_model_pt(self, deploy_params: dict):
+        """
+        Downloads the PyTorch checkpoint from Team Files if TensorRT is failed to load.
+        """
+        team_id = sly_env.team_id()
+
+        checkpoint_name = sly_fs.get_file_name(deploy_params["model_files"]["checkpoint"])
+        artifacts_dir = deploy_params["model_info"]["artifacts_dir"]
+        checkpoints_dir = os.path.join(artifacts_dir, "checkpoints")
+        checkpoint_ext = sly_fs.get_file_ext(deploy_params["model_info"]["checkpoints"][0])
+
+        pt_checkpoint_name = f"{checkpoint_name}{checkpoint_ext}"
+        remote_checkpoint_path = os.path.join(checkpoints_dir, pt_checkpoint_name)
+        local_checkpoint_path = os.path.join(self.model_dir, pt_checkpoint_name)
+
+        file_info = self.api.file.get_info_by_path(team_id, remote_checkpoint_path)
+        file_size = file_info.sizeb
+        if self.gui is not None:
+            with self.gui.download_progress(
+                message=f"Fallback. Downloading PyTorch checkpoint: '{pt_checkpoint_name}'",
+                total=file_size,
+                unit="bytes",
+                unit_scale=True,
+            ) as download_pbar:
+                self.gui.download_progress.show()
+                self.api.file.download(team_id, remote_checkpoint_path, local_checkpoint_path, progress_cb=download_pbar.update)
+                self.gui.download_progress.hide()
+        else:
+            self.api.file.download(team_id, remote_checkpoint_path, local_checkpoint_path)
+
+        return local_checkpoint_path
+
+    def _remove_exported_checkpoints(self, checkpoint_path: str):
+        """
+        Removes the exported checkpoints for provided PyTorch checkpoint path.
+        """
+        checkpoint_ext = sly_fs.get_file_ext(checkpoint_path)
+        onnx_path = checkpoint_path.replace(checkpoint_ext, ".onnx")
+        engine_path = checkpoint_path.replace(checkpoint_ext, ".engine")
+        if os.path.exists(onnx_path):
+            sly_fs.silent_remove(onnx_path)
+        if os.path.exists(engine_path):
+            sly_fs.silent_remove(engine_path)
+
     def _download_custom_model(self, model_files: dict, log_progress: bool = True):
         """
         Downloads the custom model data.
@@ -1060,7 +1104,35 @@ class Inference:
         self.runtime = deploy_params.get("runtime", RuntimeType.PYTORCH)
         self.model_precision = deploy_params.get("model_precision", ModelPrecision.FP32)
         self._hardware = get_hardware_info(self.device)
-        self.load_model(**deploy_params)
+
+        checkpoint_path = deploy_params["model_files"]["checkpoint"]
+        checkpoint_ext = sly_fs.get_file_ext(checkpoint_path)
+        if self.runtime == RuntimeType.TENSORRT and checkpoint_ext == ".engine":
+            try:
+                self.load_model(**deploy_params)
+            except Exception as e:
+                logger.warning(f"Failed to load model with TensorRT. Downloading PyTorch to export to TensorRT. Error: {repr(e)}")
+                checkpoint_path = self._fallback_download_custom_model_pt(deploy_params)
+                deploy_params["model_files"]["checkpoint"] = checkpoint_path
+                logger.info("Exporting PyTorch model to TensorRT...")
+                self._remove_exported_checkpoints(checkpoint_path)
+                checkpoint_path = self.export_tensorrt(deploy_params)
+                deploy_params["model_files"]["checkpoint"] = checkpoint_path
+                self.load_model(**deploy_params)
+        if checkpoint_ext in (".pt", ".pth") and not self.runtime == RuntimeType.PYTORCH:
+            if self.runtime == RuntimeType.ONNXRUNTIME:
+                logger.info("Exporting PyTorch model to ONNX...")
+                self._remove_exported_checkpoints(checkpoint_path)
+                checkpoint_path = self.export_onnx(deploy_params)
+            elif self.runtime == RuntimeType.TENSORRT:
+                logger.info("Exporting PyTorch model to TensorRT...")
+                self._remove_exported_checkpoints(checkpoint_path)
+                checkpoint_path = self.export_tensorrt(deploy_params)
+            deploy_params["model_files"]["checkpoint"] = checkpoint_path
+            self.load_model(**deploy_params)
+        else:
+            self.load_model(**deploy_params)
+
         self._model_served = True
         self._deploy_params = deploy_params
         if self._task_id is not None and is_production():
@@ -1181,6 +1253,7 @@ class Inference:
         if self._model_meta is None:
             self._set_model_meta_from_classes()
 
+
     def _set_model_meta_custom_model(self, model_info: dict):
         model_meta = model_info.get("model_meta")
         if model_meta is None:
@@ -2526,7 +2599,6 @@ class Inference:
         timer.daemon = True
         timer.start()
         self._freeze_timer = timer
-        logger.debug("Model will be frozen in %s seconds due to inactivity.", self._inactivity_timeout)
 
     def _set_served_callback(self):
         self._model_served = True
@@ -4214,6 +4286,11 @@ class Inference:
             return
         self.gui.model_source_tabs.set_active_tab(ModelSource.PRETRAINED)
 
+    def export_onnx(self, deploy_params: dict):
+        raise NotImplementedError("Have to be implemented in child class after inheritance")
+
+    def export_tensorrt(self, deploy_params: dict):
+        raise NotImplementedError("Have to be implemented in child class after inheritance")
 
 def _exclude_duplicated_predictions(
     api: Api,

supervisely/nn/tracker/visualize.py

@@ -161,8 +161,8 @@ class TrackingVisualizer:
         process = (
             ffmpeg
             .input(str(video_path))
-            .output('pipe:', format='rawvideo', pix_fmt='bgr24')
-            .run_async(pipe_stdout=True, pipe_stderr=True)
+            .output('pipe:', format='rawvideo', pix_fmt='bgr24', loglevel='quiet')
+            .run_async(pipe_stdout=True, pipe_stderr=False)
         )
         
         try:
@@ -177,7 +177,10 @@ class TrackingVisualizer:
                 frame = np.frombuffer(frame_data, np.uint8).reshape([height, width, 3])
                 yield frame_idx, frame
                 frame_idx += 1
-                
+
+        except ffmpeg.Error as e:
+            logger.error(f"ffmpeg error: {e.stderr.decode() if e.stderr else str(e)}", exc_info=True)
+
         finally:
             process.stdout.close()
             if process.stderr:

{supervisely-6.73.442.dist-info → supervisely-6.73.443.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: supervisely
-Version: 6.73.442
+Version: 6.73.443
 Summary: Supervisely Python SDK.
 Home-page: https://github.com/supervisely/supervisely
 Author: Supervisely

{supervisely-6.73.442.dist-info → supervisely-6.73.443.dist-info}/RECORD

@@ -904,7 +904,7 @@ supervisely/nn/benchmark/visualization/widgets/table/__init__.py,sha256=47DEQpj8
 supervisely/nn/benchmark/visualization/widgets/table/table.py,sha256=atmDnF1Af6qLQBUjLhK18RMDKAYlxnsuVHMSEa5a-e8,4319
 supervisely/nn/inference/__init__.py,sha256=QFukX2ip-U7263aEPCF_UCFwj6EujbMnsgrXp5Bbt8I,1623
 supervisely/nn/inference/cache.py,sha256=rfmb1teJ9lNDfisUSh6bwDCVkPZocn8GMvDgLQktnbo,35023
-supervisely/nn/inference/inference.py,sha256=sr70lji4u7V5MsZBiUBBUuc8_dL_FsNLboAPexAq0HU,203145
+supervisely/nn/inference/inference.py,sha256=aFFFuNhRV8m4Ch4eB-zvw1gnhz7X9WJZz-bQ3g6wUyM,207124
 supervisely/nn/inference/inference_request.py,sha256=y6yw0vbaRRcEBS27nq3y0sL6Gmq2qLA_Bm0GrnJGegE,14267
 supervisely/nn/inference/session.py,sha256=XUqJ_CqHk3ZJYkWxdnErN_6afCpIBU76nq6Ek7DiOQI,35792
 supervisely/nn/inference/uploader.py,sha256=Dn5MfMRq7tclEWpP0B9fJjTiQPBpwumfXxC8-lOYgnM,5659
@@ -999,7 +999,7 @@ supervisely/nn/tracker/base_tracker.py,sha256=2d23JlHizOqVye324YT20EE8RP52uwoQUk
 supervisely/nn/tracker/botsort_tracker.py,sha256=F2OaoeK1EAlBKAY95Fd9ZooZIlOZBh4YThhzmKNyP6w,10224
 supervisely/nn/tracker/calculate_metrics.py,sha256=JjXI4VYWYSZ5j2Ed81FNYozkS3v2UAM73ztjLrHGg58,10434
 supervisely/nn/tracker/utils.py,sha256=4WdtFHSEx5Buq6aND7cD21FdH0sGg92QNIMTpC9EGD4,10126
-supervisely/nn/tracker/visualize.py,sha256=DY2cnRm4w6_e47xVuF9fwSnOP27ZWTRfv4MFPCyxsD4,20146
+supervisely/nn/tracker/visualize.py,sha256=s0QdqEE0rcKsXow3b9lSoM9FMtms-p26HqsfNiZ2W_s,20286
 supervisely/nn/tracker/botsort/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 supervisely/nn/tracker/botsort/botsort_config.yaml,sha256=q_7Gp1-15lGYOLv7JvxVJ69mm6hbCLbUAl_ZBOYNGpw,535
 supervisely/nn/tracker/botsort/osnet_reid/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1127,9 +1127,9 @@ supervisely/worker_proto/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 supervisely/worker_proto/worker_api_pb2.py,sha256=VQfi5JRBHs2pFCK1snec3JECgGnua3Xjqw_-b3aFxuM,59142
 supervisely/worker_proto/worker_api_pb2_grpc.py,sha256=3BwQXOaP9qpdi0Dt9EKG--Lm8KGN0C5AgmUfRv77_Jk,28940
 supervisely_lib/__init__.py,sha256=yRwzEQmVwSd6lUQoAUdBngKEOlnoQ6hA9ZcoZGJRNC4,331
-supervisely-6.73.442.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-supervisely-6.73.442.dist-info/METADATA,sha256=QnoEurAbF56kquhG3igJuoDl4p4SO571wv2ngekr6Fk,35480
-supervisely-6.73.442.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-supervisely-6.73.442.dist-info/entry_points.txt,sha256=U96-5Hxrp2ApRjnCoUiUhWMqijqh8zLR03sEhWtAcms,102
-supervisely-6.73.442.dist-info/top_level.txt,sha256=kcFVwb7SXtfqZifrZaSE3owHExX4gcNYe7Q2uoby084,28
-supervisely-6.73.442.dist-info/RECORD,,
+supervisely-6.73.443.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+supervisely-6.73.443.dist-info/METADATA,sha256=XiNJG5fGmMFlBABN2mptuO55-VvjCbOxLMflB-PHDOE,35480
+supervisely-6.73.443.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+supervisely-6.73.443.dist-info/entry_points.txt,sha256=U96-5Hxrp2ApRjnCoUiUhWMqijqh8zLR03sEhWtAcms,102
+supervisely-6.73.443.dist-info/top_level.txt,sha256=kcFVwb7SXtfqZifrZaSE3owHExX4gcNYe7Q2uoby084,28
+supervisely-6.73.443.dist-info/RECORD,,