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

supervisely/nn/tracking/boxmot.py

@@ -0,0 +1,114 @@
+from typing import Iterable, List, Optional, Union
+
+import numpy as np
+
+from supervisely.annotation.annotation import Annotation
+from supervisely.geometry.rectangle import Rectangle
+from supervisely.nn.model.model_api import Prediction
+from supervisely.video_annotation.frame import Frame
+from supervisely.video_annotation.frame_collection import FrameCollection
+from supervisely.video_annotation.video_annotation import VideoAnnotation
+from supervisely.video_annotation.video_figure import VideoFigure
+from supervisely.video_annotation.video_object import VideoObject
+from supervisely.video_annotation.video_object_collection import VideoObjectCollection
+
+
+def _none_generator():
+    while True:
+        yield None
+
+
+def apply_boxmot(
+    tracker,
+    predictions: Union[List[Prediction], List[Annotation]],
+    class_names: List[str],
+    frames: Optional[Iterable[np.ndarray]] = None,
+) -> VideoAnnotation:
+    if frames is None:
+        frames = _none_generator()
+    results = []
+    annotations = []
+    frames_count = 0
+    for prediction, frame in zip(predictions, frames):
+        frames_count += 1
+        if isinstance(prediction, Prediction):
+            annotation = prediction.annotation
+            if frame is None:
+                frame = prediction.load_image()
+        else:
+            annotation = prediction
+        frame_shape = frame.shape[:2]
+        annotations.append(annotation)
+        detections = to_boxes(annotation, class_names)  # N x (x, y, x, y, conf, label)
+        tracks = tracker.update(
+            detections, frame
+        )  # M x (x, y, x, y, track_id, conf, label, det_id)
+        results.append(tracks)
+    return create_video_annotation(annotations, results, class_names, frame_shape, frames_count)
+
+
+def to_boxes(ann: Annotation, class_names: List[str]) -> np.ndarray:
+    """
+    Convert annotation to detections array in boxmot format.
+    :param ann: Supervisely Annotation object
+    :type ann: Annotation
+    :param class_names: model class names
+    :type class_names: List[str]
+    :return: detections array N x (x, y, x, y, conf, label)
+    :rtype: np.ndarray
+    """
+    # convert ann to N x (x, y, x, y, conf, cls) np.array
+    cls2label = {class_name: i for i, class_name in enumerate(class_names)}
+    detections = []
+    for label in ann.labels:
+        cat = cls2label[label.obj_class.name]
+        bbox = label.geometry.to_bbox()
+        conf = label.tags.get("confidence").value
+        detections.append([bbox.left, bbox.top, bbox.right, bbox.bottom, conf, cat])
+    detections = np.array(detections)
+    return detections
+
+
+def create_video_annotation(
+    annotations: List[Annotation],
+    tracking_results: list,
+    class_names: List[str],
+    frame_shape: tuple,
+    frames_count: int,
+) -> VideoAnnotation:
+    img_h, img_w = frame_shape
+    video_objects = {}  # track_id -> VideoObject
+    frames = []
+    cat2obj = {}
+    name2cat = {class_name: i for i, class_name in enumerate(class_names)}
+    obj_classes = {}
+    for annotation in annotations:
+        for label in annotation.labels:
+            obj_classes.setdefault(label.obj_class.name, label.obj_class)
+    for obj_name, cat in name2cat.items():
+        obj_class = obj_classes.get(obj_name)
+        if obj_class is None:
+            raise ValueError(f"Object class {obj_name} not found in annotations.")
+        cat2obj[cat] = obj_class
+    for i, tracks in enumerate(tracking_results):
+        frame_figures = []
+        for track in tracks:
+            # crop bbox to image size
+            dims = np.array([img_w, img_h, img_w, img_h]) - 1
+            track[:4] = np.clip(track[:4], 0, dims)
+            x1, y1, x2, y2, track_id, conf, cat = track[:7]
+            cat = int(cat)
+            track_id = int(track_id)
+            rect = Rectangle(y1, x1, y2, x2)
+            video_object = video_objects.setdefault(track_id, VideoObject(cat2obj[cat]))
+            frame_figures.append(VideoFigure(video_object, rect, i))
+        frames.append(Frame(i, frame_figures))
+
+    objects = list(video_objects.values())
+    video_annotation = VideoAnnotation(
+        img_size=frame_shape,
+        frames_count=frames_count,
+        objects=VideoObjectCollection(objects),
+        frames=FrameCollection(frames),
+    )
+    return video_annotation

supervisely/nn/tracking/tracking.py

@@ -0,0 +1,24 @@
+from typing import Any, Callable
+
+from supervisely.nn.model.model_api import ModelAPI
+from supervisely.nn.tracking.boxmot import apply_boxmot
+from supervisely.video_annotation.video_annotation import VideoAnnotation
+
+
+def _get_apply_fn(tracker: Any) -> Callable:
+    if tracker.__class__.__module__.startswith("boxmot"):
+        return apply_boxmot
+    else:
+        raise ValueError(
+            f"Tracker {tracker.__class__.__module__} is not supported. Please, use boxmot tracker."
+        )
+
+
+def track(video_id: int, tracker, detector: ModelAPI, **kwargs) -> VideoAnnotation:
+    apply_fn = _get_apply_fn(tracker)
+    if "classes" in kwargs:
+        classes = kwargs["classes"]
+    else:
+        classes = detector.get_classes()
+    predictions = detector.predict_detached(video_id=video_id, **kwargs)
+    return apply_fn(tracker, predictions, classes)

supervisely/nn/training/train_app.py

@@ -525,19 +525,19 @@ class TrainApp:
         if not success:
             raise ValueError(f"{reason}. Failed to upload artifacts")
 
-        # Step 3. Preprocess artifacts
-        experiment_info = self._preprocess_artifacts(experiment_info)
+        # Step 3. Create model meta according to model CV task type
+        model_meta = self.create_model_meta(experiment_info["task_type"])
+
+        # Step 4. Preprocess artifacts
+        experiment_info = self._preprocess_artifacts(experiment_info, model_meta)
 
-        # Step 4. Postprocess splits
+        # Step 5. Postprocess splits
         train_splits_data = self._postprocess_splits()
 
-        # Step 5. Upload artifacts
+        # Step 6. Upload artifacts
         self._set_text_status("uploading")
         remote_dir, file_info = self._upload_artifacts()
 
-        # Step 6. Create model meta according to model CV task type
-        model_meta = self.create_model_meta(experiment_info["task_type"])
-
         # Step 7. [Optional] Run Model Benchmark
         mb_eval_lnk_file_info, mb_eval_report = None, None
         mb_eval_report_id, eval_metrics = None, {}
@@ -579,8 +579,10 @@ class TrainApp:
         export_weights = {}
         if self.gui.hyperparameters_selector.is_export_required():
             try:
-                export_weights = self._export_weights(experiment_info)
-                export_weights = self._upload_export_weights(export_weights, remote_dir)
+                export_weights, export_classes_path = self._export_weights(experiment_info)
+                export_weights = self._upload_export_weights(
+                    export_weights, export_classes_path, remote_dir
+                )
             except Exception as e:
                 logger.error(f"Export weights failed: {e}")
 
@@ -1345,7 +1347,7 @@ class TrainApp:
         }
         return splits_data
 
-    def _preprocess_artifacts(self, experiment_info: dict) -> None:
+    def _preprocess_artifacts(self, experiment_info: dict, model_meta: ProjectMeta) -> None:
         """
         Preprocesses and move the artifacts generated by the training process to output directories.
 
@@ -1397,12 +1399,42 @@ class TrainApp:
 
         new_checkpoint_paths = []
         best_checkpoints_name = experiment_info["best_checkpoint"]
+
+        # Prepare model files
+        try:
+            model_files = {}
+            for file in experiment_info["model_files"]:
+                with open(experiment_info["model_files"][file], "r") as f:
+                    model_files[file] = f.read()
+        except Exception as e:
+            logger.warning(f"Error loading model files: {e}")
+            model_files = {}
+
         for checkpoint_path in checkpoint_paths:
-            new_checkpoint_path = join(
-                self._output_checkpoints_dir,
-                sly_fs.get_file_name_with_ext(checkpoint_path),
-            )
+            checkpoint_name = sly_fs.get_file_name_with_ext(checkpoint_path)
+            new_checkpoint_path = join(self._output_checkpoints_dir, checkpoint_name)
             shutil.move(checkpoint_path, new_checkpoint_path)
+            if len(model_files) > 0:
+                try:
+                    # pylint: disable=import-error
+                    import torch
+
+                    state_dict = torch.load(new_checkpoint_path)
+                    state_dict["model_info"] = {
+                        "model_name": experiment_info["model_name"],
+                        "framework": self.framework_name,
+                        "checkpoint": checkpoint_name,
+                        "experiment": self.gui.training_process.get_experiment_name(),
+                    }
+                    state_dict["model_meta"] = model_meta.to_json()
+                    state_dict["model_files"] = model_files
+                    torch.save(state_dict, new_checkpoint_path)
+                except Exception as e:
+                    logger.warning(
+                        f"Error writing info to checkpoint: '{checkpoint_name}'. Error:{e}"
+                    )
+                    continue
+
             new_checkpoint_paths.append(new_checkpoint_path)
             if sly_fs.get_file_name_with_ext(checkpoint_path) == best_checkpoints_name:
                 experiment_info["best_checkpoint"] = new_checkpoint_path
@@ -2543,17 +2575,27 @@ class TrainApp:
                     export_weights[RuntimeType.TENSORRT] = tensorrt_path
             except Exception as e:
                 logger.error(f"Failed to export TensorRT model: {e}")
-                
-        return export_weights
+
+        export_classes_path = None
+        if len(export_weights) > 0:
+            export_classes = {idx: cls_name for idx, cls_name in enumerate(self.classes)}
+            export_dir = join(self.output_dir, self._export_dir_name)
+            sly_fs.mkdir(export_dir)
+            export_classes_path = join(export_dir, "classes.json")
+            sly_json.dump_json_file(export_classes, export_classes_path)
+        return export_weights, export_classes_path
 
     def _upload_export_weights(
-        self, export_weights: Dict[str, str], remote_dir: str
+        self, export_weights: Dict[str, str], export_classes_path: str, remote_dir: str
     ) -> Dict[str, str]:
         """Uploads export weights (any other specified formats) to Supervisely Team Files.
         The default export is handled by the `_upload_artifacts` method."""
         file_dest_paths = []
         size = 0
-        for path in export_weights.values():
+        files_to_upload = list(export_weights.values())
+        if export_classes_path is not None:
+            files_to_upload.append(export_classes_path)
+        for path in files_to_upload:
             file_name = sly_fs.get_file_name_with_ext(path)
             file_dest_paths.append(join(remote_dir, self._export_dir_name, file_name))
             size += sly_fs.get_file_size(path)
@@ -2568,7 +2610,7 @@ class TrainApp:
             self.progress_bar_main.show()
             self._api.file.upload_bulk_fast(
                 team_id=self.team_id,
-                src_paths=export_weights.values(),
+                src_paths=files_to_upload,
                 dst_paths=file_dest_paths,
                 progress_cb=export_upload_main_pbar.update,
             )

supervisely/nn/utils.py

@@ -1,8 +1,6 @@
 from dataclasses import dataclass
-from typing import List
 
-from supervisely.nn.artifacts.artifacts import TrainInfo
-from supervisely.nn.experiments import ExperimentInfo
+import psutil
 
 
 class ModelSource:
@@ -61,6 +59,48 @@ def _get_model_name(model_info: dict):
     name = model_info.get("model_name")
     if not name:
         name = model_info.get("meta", {}).get("model_name")
+    if not name:
+        name = model_info.get("Model")
     if not name:
         raise ValueError("Model name not found not in model_info nor in meta.")
     return name
+
+
+def get_ram_usage():
+    memory = psutil.virtual_memory()
+    return memory.used, memory.total
+
+
+def get_gpu_usage(device: str = None):
+    if device == "cpu":
+        return None, None
+    try:
+        import torch
+    except Exception as e:
+        from supervisely import logger
+        logger.warning(f"Cannot import torch. Install PyTorch to get GPU usage info. Error: {e}")
+        return None, None
+    if not torch.cuda.is_available():
+        return None, None
+    gpu_index = None
+    if device is None or device in ["", "auto", "cuda"]:
+        gpu_index = torch.cuda.current_device()
+    elif isinstance(device, int):
+        gpu_index = device
+    elif device.startswith("cuda:"):
+        try:
+            gpu_index = int(device.split(":")[-1])
+        except ValueError:
+            return None, None
+    else:
+        for i in range(torch.cuda.device_count()):
+            if device == torch.cuda.get_device_name(i):
+                gpu_index = i
+                break
+        if gpu_index is None:
+            return None, None
+    if gpu_index is None or gpu_index > torch.cuda.device_count() or gpu_index < 0:
+        return None, None
+    allocated = torch.cuda.memory_allocated(gpu_index)
+    total = torch.cuda.get_device_properties(gpu_index).total_memory
+    return allocated, total

supervisely/task/progress.py

@@ -79,6 +79,8 @@ class Progress:
         is_size: Optional[bool] = False,
         need_info_log: Optional[bool] = False,
         min_report_percent: Optional[int] = 1,
+        log_extra: Optional[Dict[str, str]] = None,
+        update_task_progress: Optional[bool] = True,
     ):
         self.is_size = is_size
         self.message = message
@@ -94,6 +96,8 @@ class Progress:
         self.logger = logger if ext_logger is None else ext_logger
         self.report_every = max(1, math.ceil((total_cnt or 0) / 100 * min_report_percent))
         self.need_info_log = need_info_log
+        self.log_extra = log_extra
+        self.update_task_progress = update_task_progress
 
         mb5 = 5 * 1024 * 1024
         if self.is_size and self.is_total_unknown:
@@ -169,9 +173,15 @@ class Progress:
             extra["current_label"] = self.current_label
             extra["total_label"] = self.total_label
 
-        self.logger.info("progress", extra=extra)
+        if self.log_extra:
+            extra.update(self.log_extra)
+
+        if self.update_task_progress:
+            self.logger.info("progress", extra=extra)
         if self.need_info_log is True:
-            self.logger.info(f"{self.message} [{self.current_label} / {self.total_label}]")
+            self.logger.info(
+                f"{self.message} [{self.current_label} / {self.total_label}]", extra=self.log_extra
+            )
 
     def need_report(self) -> bool:
         if (

supervisely/video/video.py

@@ -4,7 +4,10 @@
 from __future__ import annotations
 
 import os
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, Generator, List, Optional, Tuple
+
+import cv2
+import numpy as np
 
 from supervisely import logger as default_logger
 from supervisely._utils import abs_url, is_development, rand_str
@@ -519,3 +522,106 @@ def get_labeling_tool_link(url: str, name: Optional[str] = "open in labeling too
     :rtype: str
     """
     return f'<a href="{url}" rel="noopener noreferrer" target="_blank">{name}<i class="zmdi zmdi-open-in-new" style="margin-left: 5px"></i></a>'
+
+
+class VideoFrameReader:
+    def __init__(self, video_path: str, frame_indexes: List[int] = None):
+        self.video_path = video_path
+        self.frame_indexes = frame_indexes
+        self.vr = None
+        self.cap = None
+        self.prev_idx = -1
+
+    def _ensure_initialized(self):
+        if self.vr is None and self.cap is None:
+            try:
+                import decord
+
+                self.vr = decord.VideoReader(str(self.video_path))
+            except ImportError:
+                default_logger.debug(
+                    "Decord is not installed. Falling back to OpenCV for video reading."
+                )
+                self.cap = cv2.VideoCapture(str(self.video_path))
+
+    def close(self):
+        if self.vr is not None:
+            self.vr = None
+        if self.cap is not None:
+            self.cap.release()
+            self.cap = None
+        self.prev_idx = -1
+
+    def __enter__(self):
+        self._ensure_initialized()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def __del__(self):
+        self.close()
+
+    def iterate_frames(self, frame_indexes: List[int] = None) -> Generator[np.ndarray, None, None]:
+        self._ensure_initialized()
+        if frame_indexes is None:
+            frame_indexes = self.frame_indexes
+        if self.vr is not None:
+            if frame_indexes is None:
+                frame_indexes = range(len(self.vr))
+            for frame_index in frame_indexes:
+                frame = self.vr[frame_index]
+                yield frame.asnumpy()
+        else:
+            if frame_indexes is None:
+                frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
+                frame_indexes = range(frame_count)
+            for frame_index in frame_indexes:
+                if 1 > frame_index - self.prev_idx < 20:
+                    while self.prev_idx < frame_index - 1:
+                        self.cap.read()
+                if frame_index != self.prev_idx + 1:
+                    self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
+                ret, frame = self.cap.read()
+                if not ret:
+                    raise KeyError(f"Frame {frame_index} not found in video {self.video_path}")
+                yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                self.prev_idx = frame_index
+
+    def read_frames(self, frame_indexes: List[int] = None) -> List[np.ndarray]:
+        return list(self.iterate_frames(frame_indexes))
+
+    def __iter__(self):
+        return self.iterate_frames()
+
+    def __next__(self):
+        if not hasattr(self, "_frame_generator"):
+            self._frame_generator = self.iterate_frames()
+        try:
+            return next(self._frame_generator)
+        except StopIteration:
+            self._frame_generator = None
+            raise
+
+    def frame_size(self):
+        self._ensure_initialized()
+        if self.vr is not None:
+            return self.vr[0].shape[:2]
+        else:
+            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            return height, width
+
+    def frames_count(self):
+        self._ensure_initialized()
+        if self.vr is not None:
+            return len(self.vr)
+        else:
+            return int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    def fps(self):
+        self._ensure_initialized()
+        if self.vr is not None:
+            return self.vr.get_avg_fps()
+        else:
+            return int(self.cap.get(cv2.CAP_PROP_FPS))

{supervisely-6.73.357.dist-info → supervisely-6.73.358.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: supervisely
-Version: 6.73.357
+Version: 6.73.358
 Summary: Supervisely Python SDK.
 Home-page: https://github.com/supervisely/supervisely
 Author: Supervisely
@@ -128,6 +128,7 @@ Requires-Dist: scikit-learn; extra == "tracking"
 Requires-Dist: faiss-gpu; extra == "tracking"
 Requires-Dist: tabulate; extra == "tracking"
 Requires-Dist: tensorboard; extra == "tracking"
+Requires-Dist: decord; extra == "tracking"
 Provides-Extra: training
 Requires-Dist: pycocotools; extra == "training"
 Requires-Dist: scikit-learn; extra == "training"