supervisely 6.73.452__py3-none-any.whl → 6.73.513__py3-none-any.whl

supervisely/nn/tracker/visualize.py

@@ -1,21 +1,21 @@
-from typing import Union, Dict, List, Tuple, Iterator, Optional
-import numpy as np
+import shutil
+import tempfile
+from collections import defaultdict
+from pathlib import Path
+from typing import Iterator, List, Optional, Tuple, Union
+
 import cv2
 import ffmpeg
-from pathlib import Path
-from collections import defaultdict
-from dataclasses import dataclass
-import tempfile
-import shutil
+import numpy as np
 
-import supervisely as sly
-from supervisely import logger
+from supervisely import VideoAnnotation, logger
+from supervisely.geometry.geometry import Geometry
 from supervisely.nn.model.prediction import Prediction
-from supervisely import VideoAnnotation
 from supervisely.nn.tracker.utils import predictions_to_video_annotation
 
 
 class TrackingVisualizer:
+
     def __init__(
         self,
         show_labels: bool = True,
@@ -29,7 +29,7 @@ class TrackingVisualizer:
         codec: str = "mp4",
         output_fps: float = 30.0,
         colorize_tracks: bool = True,
-
+        trajectory_thickness: int = 2,
     ):
         """
         Initialize the visualizer with configuration.
@@ -58,6 +58,7 @@ class TrackingVisualizer:
         self.text_scale = text_scale
         self.text_thickness = text_thickness
         self.trajectory_length = trajectory_length
+        self.trajectory_thickness = trajectory_thickness
         self.colorize_tracks = colorize_tracks
 
         # Output settings
@@ -71,7 +72,7 @@ class TrackingVisualizer:
         self.track_colors = {}
         self.color_palette = self._generate_color_palette()
         self._temp_dir = None
-        
+
     def _generate_color_palette(self, num_colors: int = 100) -> List[Tuple[int, int, int]]:
         """
         Generate bright, distinct color palette for track visualization.
@@ -88,11 +89,11 @@ class TrackingVisualizer:
             bgr_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)[0][0]
             colors.append(tuple(map(int, bgr_color)))
         return colors
-    
+
     def _get_track_color(self, track_id: int) -> Tuple[int, int, int]:
         """Get consistent color for track ID from palette."""
         return self.color_palette[track_id % len(self.color_palette)]
-    
+
     def _get_video_info(self, video_path: Path) -> Tuple[int, int, float, int]:
         """
         Get video metadata using ffmpeg.
@@ -104,13 +105,13 @@ class TrackingVisualizer:
             probe = ffmpeg.probe(str(video_path))
             video_stream = next((stream for stream in probe['streams'] 
                                if stream['codec_type'] == 'video'), None)
-            
+
             if video_stream is None:
                 raise ValueError(f"No video stream found in: {video_path}")
-            
+
             width = int(video_stream['width'])
             height = int(video_stream['height'])
-            
+
             # Extract FPS
             fps_str = video_stream.get('r_frame_rate', '30/1')
             if '/' in fps_str:
@@ -118,19 +119,19 @@ class TrackingVisualizer:
                 fps = num / den if den != 0 else 30.0
             else:
                 fps = float(fps_str)
-            
+
             # Get total frames
             total_frames = int(video_stream.get('nb_frames', 0))
             if total_frames == 0:
                 # Fallback: estimate from duration and fps
                 duration = float(video_stream.get('duration', 0))
                 total_frames = int(duration * fps) if duration > 0 else 0
-            
+
             return width, height, fps, total_frames
-            
+
         except Exception as e:
             raise ValueError(f"Could not read video metadata {video_path}: {str(e)}")
-    
+
     def _create_frame_iterator(self, source: Union[str, Path]) -> Iterator[Tuple[int, np.ndarray]]:
         """
         Create iterator that yields (frame_index, frame) tuples.
@@ -142,38 +143,38 @@ class TrackingVisualizer:
             Tuple of (frame_index, frame_array)
         """
         source = Path(source)
-        
+
         if source.is_file():
             yield from self._iterate_video_frames(source)
         elif source.is_dir():
             yield from self._iterate_directory_frames(source)
         else:
             raise ValueError(f"Source must be a video file or directory, got: {source}")
-    
+
     def _iterate_video_frames(self, video_path: Path) -> Iterator[Tuple[int, np.ndarray]]:
         """Iterate through video frames using ffmpeg."""
         width, height, fps, total_frames = self._get_video_info(video_path)
-        
+
         # Store video info for later use
         self.source_fps = fps
         self.frame_size = (width, height)
-        
+
         process = (
             ffmpeg
             .input(str(video_path))
             .output('pipe:', format='rawvideo', pix_fmt='bgr24', loglevel='quiet')
             .run_async(pipe_stdout=True, pipe_stderr=False)
         )
-        
+
         try:
             frame_size_bytes = width * height * 3
             frame_idx = 0
-            
+
             while True:
                 frame_data = process.stdout.read(frame_size_bytes)
                 if len(frame_data) != frame_size_bytes:
                     break
-                
+
                 frame = np.frombuffer(frame_data, np.uint8).reshape([height, width, 3])
                 yield frame_idx, frame
                 frame_idx += 1
@@ -186,26 +187,26 @@ class TrackingVisualizer:
             if process.stderr:
                 process.stderr.close()
             process.wait()
-    
+
     def _iterate_directory_frames(self, frames_dir: Path) -> Iterator[Tuple[int, np.ndarray]]:
         """Iterate through image frames in directory."""
         if not frames_dir.is_dir():
             raise ValueError(f"Directory does not exist: {frames_dir}")
-        
+
         # Support common image extensions
         extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
         image_files = []
         for ext in extensions:
             image_files.extend(frames_dir.glob(f'*{ext}'))
             image_files.extend(frames_dir.glob(f'*{ext.upper()}'))
-        
+
         image_files = sorted(image_files)
         if not image_files:
             raise ValueError(f"No image files found in directory: {frames_dir}")
-        
+
         # Set fps from config for image sequences
         self.source_fps = self.output_fps
-        
+
         for frame_idx, img_path in enumerate(image_files):
             frame = cv2.imread(str(img_path))
             if frame is not None:
@@ -215,7 +216,7 @@ class TrackingVisualizer:
                 yield frame_idx, frame
             else:
                 logger.warning(f"Could not read image: {img_path}")
-    
+
     def _extract_tracks_from_annotation(self) -> None:
         """
         Extract tracking data from Supervisely VideoAnnotation.
@@ -224,29 +225,22 @@ class TrackingVisualizer:
         """
         self.tracks_by_frame = defaultdict(list)
         self.track_colors = {}
-        
+
         # Map object keys to track info
         objects = {}
         for i, obj in enumerate(self.annotation.objects):
             objects[obj.key] = (i, obj.obj_class.name)
-        
+
         # Extract tracks from frames
         for frame in self.annotation.frames:
             frame_idx = frame.index
             for figure in frame.figures:
-                if figure.geometry.geometry_name() != 'rectangle':
-                    continue
-                    
                 object_key = figure.parent_object.key
                 if object_key not in objects:
                     continue
-                    
+
                 track_id, class_name = objects[object_key]
-                
-                # Extract bbox coordinates
-                rect = figure.geometry
-                bbox = (rect.left, rect.top, rect.right, rect.bottom)
-                
+
                 if track_id not in self.track_colors:
                     if self.colorize_tracks:
                         # auto-color override everything
@@ -263,26 +257,30 @@ class TrackingVisualizer:
 
                     self.track_colors[track_id] = color
 
-                
-                self.tracks_by_frame[frame_idx].append((track_id, bbox, class_name))
-        
+                self.tracks_by_frame[frame_idx].append((track_id, figure.geometry, class_name))
+
         logger.info(f"Extracted tracks from {len(self.tracks_by_frame)} frames")
-        
-    def _draw_detection(self, img: np.ndarray, track_id: int, bbox: Tuple[int, int, int, int], 
-                    class_name: str) -> Optional[Tuple[int, int]]:
+
+    def _draw_detection(
+        self, img: np.ndarray, track_id: int, geometry: Geometry, class_name: str
+    ) -> Optional[Tuple[int, int]]:
         """
         Draw single detection with track ID and class label.
         Returns the center point of the bbox for trajectory drawing.
         """
+        rect = geometry.to_bbox()
+        bbox = (rect.left, rect.top, rect.right, rect.bottom)
+
         x1, y1, x2, y2 = map(int, bbox)
 
         if x2 <= x1 or y2 <= y1:
             return None
-        
+
         color = self.track_colors[track_id]
 
         # Draw bounding box
-        cv2.rectangle(img, (x1, y1), (x2, y2), color, self.box_thickness)
+        geometry.draw_contour(img, color=color, thickness=self.box_thickness)
+        # cv2.rectangle(img, (x1, y1), (x2, y2), color, self.box_thickness)
 
         # Draw label if enabled
         if self.show_labels:
@@ -304,7 +302,6 @@ class TrackingVisualizer:
         # Return center point for trajectory
         return (x1 + x2) // 2, (y1 + y2) // 2
 
-
     def _draw_trajectories(self, img: np.ndarray) -> None:
         """Draw trajectory lines for all tracks, filtering out big jumps."""
         if not self.show_trajectories:
@@ -312,24 +309,24 @@ class TrackingVisualizer:
 
         max_jump = 200  
 
-        for track_id, centers in self.track_centers.items():
-            if len(centers) < 2:
+        for centers_with_colors in self.track_centers.values():
+
+            if len(centers_with_colors) < 2:
                 continue
 
-            color = self.track_colors[track_id]
-            points = centers[-self.trajectory_length:]
+            points, colors = zip(*centers_with_colors[-self.trajectory_length :])
 
             for i in range(1, len(points)):
+                color = colors[i]
                 p1, p2 = points[i - 1], points[i]
                 if p1 is None or p2 is None:
                     continue
-              
+
                 if np.hypot(p2[0] - p1[0], p2[1] - p1[1]) > max_jump:
                     continue
-                cv2.line(img, p1, p2, color, 2)
+                cv2.line(img, p1, p2, color, self.trajectory_thickness)
                 cv2.circle(img, p1, 3, color, -1)
 
-        
     def _process_single_frame(self, frame: np.ndarray, frame_idx: int) -> np.ndarray:
         """
         Process single frame: add annotations and return processed frame.
@@ -345,29 +342,30 @@ class TrackingVisualizer:
         active_ids = set()
         # Draw detections for current frame
         if frame_idx in self.tracks_by_frame:
-            for track_id, bbox, class_name in self.tracks_by_frame[frame_idx]:
-                center = self._draw_detection(img, track_id, bbox, class_name)
-                self.track_centers[track_id].append(center)
+            for track_id, geometry, class_name in self.tracks_by_frame[frame_idx]:
+                center = self._draw_detection(img, track_id, geometry, class_name)
+                color = self.track_colors[track_id]
+                self.track_centers[track_id].append((center, color))
                 if len(self.track_centers[track_id]) > self.trajectory_length:
                     self.track_centers[track_id].pop(0)
                 active_ids.add(track_id)
-        
+
         for tid in self.track_centers.keys():
             if tid not in active_ids:
-                self.track_centers[tid].append(None)
+                self.track_centers[tid].append((None, None))
                 if len(self.track_centers[tid]) > self.trajectory_length:
                     self.track_centers[tid].pop(0)
-                
+
         # Draw trajectories
         self._draw_trajectories(img)
-        
+
         # Add frame number if requested
         if self.show_frame_number:
             cv2.putText(img, f"Frame: {frame_idx + 1}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
-        
+
         return img
-    
+
     def _save_processed_frame(self, frame: np.ndarray, frame_idx: int) -> str:
         """
         Save processed frame to temporary directory.
@@ -382,7 +380,7 @@ class TrackingVisualizer:
         frame_path = self._temp_dir / f"frame_{frame_idx:08d}.jpg"
         cv2.imwrite(str(frame_path), frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
         return str(frame_path)
-    
+
     def _create_video_from_frames(self, output_path: Union[str, Path]) -> None:
         """
         Create final video from processed frames using ffmpeg.
@@ -392,10 +390,10 @@ class TrackingVisualizer:
         """
         output_path = Path(output_path)
         output_path.parent.mkdir(parents=True, exist_ok=True)
-        
+
         # Create video from frame sequence
         input_pattern = str(self._temp_dir / "frame_%08d.jpg")
-        
+
         try:
             (
                 ffmpeg
@@ -405,17 +403,17 @@ class TrackingVisualizer:
                 .run(capture_stdout=True, capture_stderr=True)
             )
             logger.info(f"Video saved to {output_path}")
-            
+
         except ffmpeg.Error as e:
             error_msg = e.stderr.decode() if e.stderr else "Unknown ffmpeg error"
             raise ValueError(f"Failed to create video: {error_msg}")
-    
+
     def _cleanup_temp_directory(self) -> None:
         """Clean up temporary directory and all its contents."""
         if self._temp_dir and self._temp_dir.exists():
             shutil.rmtree(self._temp_dir)
             self._temp_dir = None
-    
+
     def visualize_video_annotation(self, annotation: VideoAnnotation, 
                                   source: Union[str, Path], 
                                   output_path: Union[str, Path]) -> None:
@@ -433,43 +431,43 @@ class TrackingVisualizer:
         """
         if not isinstance(annotation, VideoAnnotation):
             raise TypeError(f"Annotation must be VideoAnnotation, got {type(annotation)}")
-        
+
         # Store annotation
         self.annotation = annotation
-        
+
         # Create temporary directory for processed frames
         self._temp_dir = Path(tempfile.mkdtemp(prefix="video_viz_"))
-        
+
         try:
             # Extract tracking data
             self._extract_tracks_from_annotation()
-            
+
             if not self.tracks_by_frame:
                 logger.warning("No tracking data found in annotation")
-            
+
             # Reset trajectory tracking
             self.track_centers = defaultdict(list)
-            
+
             # Process frames one by one
             frame_count = 0
             for frame_idx, frame in self._create_frame_iterator(source):
                 # Process frame
                 processed_frame = self._process_single_frame(frame, frame_idx)
-                
+
                 # Save processed frame
                 self._save_processed_frame(processed_frame, frame_idx)
-                
+
                 frame_count += 1
-                
+
                 # Progress logging
                 if frame_count % 100 == 0:
                     logger.info(f"Processed {frame_count} frames")
-            
+
             logger.info(f"Finished processing {frame_count} frames")
-            
+
             # Create final video from saved frames
             self._create_video_from_frames(output_path)
-            
+
         finally:
             # Always cleanup temporary files
             self._cleanup_temp_directory()
@@ -477,7 +475,7 @@ class TrackingVisualizer:
     def __del__(self):
         """Cleanup temporary directory on object destruction."""
         self._cleanup_temp_directory()
-        
+
 
 def visualize(
     predictions: Union[VideoAnnotation, List[Prediction]], 
@@ -519,4 +517,3 @@ def visualize(
         visualizer.visualize_video_annotation(predictions, source, output_path)
     else:
         raise TypeError(f"Predictions must be VideoAnnotation or list of Prediction, got {type(predictions)}")
-

supervisely/nn/training/gui/classes_selector.py

@@ -51,7 +51,19 @@ class ClassesSelector:
             text=f"<i class='zmdi zmdi-chart-donut' style='color: #7f858e'></i> <a href='{qa_stats_link}' target='_blank'> <b> QA & Stats </b></a>"
         )
 
-        self.classes_table = ClassesTable(project_id=project_id)
+        models = model_selector.models
+        task_types = [model["meta"]["task_type"] for model in models]
+        task_types = list(set(task_types))
+        allowed_types = []
+        for task_type in task_types:
+            if task_type.endswith("detection"):
+                allowed_types.append(Rectangle)
+            elif task_type.endswith("segmentation"):
+                allowed_types.extend([Bitmap, Polygon])
+            elif task_type == TaskType.POSE_ESTIMATION:
+                allowed_types.append(GraphNodes)
+
+        self.classes_table = ClassesTable(project_id=project_id, allowed_types=allowed_types)
         if len(classes) > 0:
             self.classes_table.select_classes(classes)
         else:
@@ -107,6 +119,9 @@ class ClassesSelector:
             TaskType.INSTANCE_SEGMENTATION: {Bitmap},
             TaskType.SEMANTIC_SEGMENTATION: {Bitmap},
             TaskType.POSE_ESTIMATION: {GraphNodes},
+            TaskType.PROMPTABLE_SEGMENTATION: {Bitmap},
+            TaskType.INTERACTIVE_SEGMENTATION: {Bitmap},
+            TaskType.PROMPT_BASED_OBJECT_DETECTION: {Rectangle},
         }
 
         if task_type not in allowed_shapes:

supervisely/nn/training/train_app.py

@@ -43,6 +43,7 @@ from supervisely import (
     logger,
 )
 from supervisely._utils import abs_url, get_filename_from_headers
+from supervisely.api.entities_collection_api import EntitiesCollectionInfo
 from supervisely.api.file_api import FileInfo
 from supervisely.app import get_synced_data_dir, show_dialog
 from supervisely.app.widgets import Progress
@@ -72,7 +73,6 @@ from supervisely.project.download import (
     is_cached,
 )
 from supervisely.template.experiment.experiment_generator import ExperimentGenerator
-from supervisely.api.entities_collection_api import EntitiesCollectionInfo
 
 
 class TrainApp:
@@ -3162,8 +3162,11 @@ class TrainApp:
 
         # Case 1: Use existing collections for training. No need to create new collections
         split_method = self.gui.train_val_splits_selector.get_split_method()
+        self.gui.train_val_splits_selector._parse_collections()
         all_train_collections = self.gui.train_val_splits_selector.all_train_collections
         all_val_collections = self.gui.train_val_splits_selector.all_val_collections
+        latest_train_collection = self.gui.train_val_splits_selector.latest_train_collection
+        latest_val_collection = self.gui.train_val_splits_selector.latest_val_collection
         if split_method == "Based on collections":
             current_selected_train_collection_ids = self.gui.train_val_splits_selector.train_val_splits.get_train_collections_ids()
             train_match = _check_match(current_selected_train_collection_ids, all_train_collections)
@@ -3178,9 +3181,6 @@ class TrainApp:
         # ------------------------------------------------------------ #
 
         # Case 2: Create new collections for selected train val splits. Need to create new collections
-        item_type = self.project_info.type
-        experiment_name = self.gui.training_process.get_experiment_name()
-
         train_collection_idx = 1
         val_collection_idx = 1
 
@@ -3191,42 +3191,41 @@ class TrainApp:
             return None
 
         # Get train collection with max idx
-        if len(all_train_collections) > 0:
-            train_indices = [_extract_index_from_col_name(collection.name, "train") for collection in all_train_collections]
-            train_indices = [idx for idx in train_indices if idx is not None]
-            if len(train_indices) > 0:
-                train_collection_idx = max(train_indices) + 1
+        if latest_train_collection:
+            train_collection_idx = (
+                _extract_index_from_col_name(latest_train_collection.name, "train") + 1
+            )
 
         # Get val collection with max idx
-        if len(all_val_collections) > 0:
-            val_indices = [_extract_index_from_col_name(collection.name, "val") for collection in all_val_collections]
-            val_indices = [idx for idx in val_indices if idx is not None]
-            if len(val_indices) > 0:
-                val_collection_idx = max(val_indices) + 1
+        if latest_val_collection:
+            val_collection_idx = _extract_index_from_col_name(latest_val_collection.name, "val") + 1
         # -------------------------------- #
 
         # Create Train Collection
         train_img_ids = list(self._train_split_item_ids)
-        train_collection_description = f"Collection with train {item_type} for experiment: {experiment_name}"
-        train_collection = self._api.entities_collection.create(self.project_id, f"train_{train_collection_idx:03d}", train_collection_description)
-        train_collection_id = getattr(train_collection, "id", None)
-        if train_collection_id is None:
-            raise AttributeError("Train EntitiesCollectionInfo object does not have 'id' attribute")
-        self._api.entities_collection.add_items(train_collection_id, train_img_ids)
-        self._train_collection_id = train_collection_id
+        self._train_collection_id = self._create_collection("train", train_collection_idx)
+        self._api.entities_collection.add_items(self._train_collection_id, train_img_ids)
 
         # Create Val Collection
         val_img_ids = list(self._val_split_item_ids)
-        val_collection_description = f"Collection with val {item_type} for experiment: {experiment_name}"
-        val_collection = self._api.entities_collection.create(self.project_id, f"val_{val_collection_idx:03d}", val_collection_description)
-        val_collection_id = getattr(val_collection, "id", None)
-        if val_collection_id is None:
-            raise AttributeError("Val EntitiesCollectionInfo object does not have 'id' attribute")
-        self._api.entities_collection.add_items(val_collection_id, val_img_ids)
-        self._val_collection_id = val_collection_id
+        self._val_collection_id = self._create_collection("val", val_collection_idx)
+        self._api.entities_collection.add_items(self._val_collection_id, val_img_ids)
 
         # Update Project Custom Data
-        self._update_project_custom_data(train_collection_id, val_collection_id)
+        self._update_project_custom_data(self._train_collection_id, self._val_collection_id)
+
+    def _create_collection(self, split_type: str, suffix: int) -> int:
+        experiment_name = self.gui.training_process.get_experiment_name()
+        description = f"Collection with {split_type} {self.project_info.type} for experiment: {experiment_name}"
+        collection = self._api.entities_collection.create(
+            project_id=self.project_id,
+            name=f"{split_type}_{suffix:03d}",
+            description=description,
+            change_name_if_conflict=True,
+        )
+        if collection is None or collection.id is None:  # pylint: disable=no-member
+            raise RuntimeError(f"Failed to create {split_type} collection")
+        return collection.id  # pylint: disable=no-member
 
     def _update_project_custom_data(self, train_collection_id: int, val_collection_id: int):
         train_info = {