supervisely 6.73.438__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,67 +143,70 @@ 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')
-            .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:
             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
-                
+
+        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:
                 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:
@@ -212,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.
@@ -221,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
@@ -260,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:
@@ -301,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:
@@ -309,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.
@@ -342,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.
@@ -379,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.
@@ -389,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
@@ -402,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:
@@ -430,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()
@@ -474,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]], 
@@ -516,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/gui/train_val_splits_selector.py

@@ -180,7 +180,13 @@ class TrainValSplitsSelector:
                 return False
 
             # Check if datasets are not empty
-            filters = [{ ApiField.FIELD: ApiField.ID, ApiField.OPERATOR: "in", ApiField.VALUE: train_dataset_id + val_dataset_id}]
+            filters = [
+                {
+                    ApiField.FIELD: ApiField.ID,
+                    ApiField.OPERATOR: "in",
+                    ApiField.VALUE: train_dataset_id + val_dataset_id,
+                }
+            ]
             selected_datasets = self.api.dataset.get_list(self.project_id, filters, recursive=True)
             datasets_count = {}
             for dataset in selected_datasets:
@@ -334,6 +340,7 @@ class TrainValSplitsSelector:
 
     def _detect_splits(self, collections_split: bool, datasets_split: bool) -> bool:
         """Detect splits based on the selected method"""
+        self._parse_collections()
         splits_found = False
         if collections_split:
             splits_found = self._detect_collections()
@@ -341,47 +348,59 @@ class TrainValSplitsSelector:
             splits_found = self._detect_datasets()
         return splits_found
 
+    def _parse_collections(self) -> None:
+        """Parse collections with train and val prefixes and set them to train_val_splits variables"""
+        all_collections = self.api.entities_collection.get_list(self.project_id)
+        existing_train_collections = [
+            collection for collection in all_collections if collection.name.startswith("train_")
+        ]
+        existing_val_collections = [
+            collection for collection in all_collections if collection.name.startswith("val_")
+        ]
+
+        self._all_train_collections = existing_train_collections
+        self._all_val_collections = existing_val_collections
+        self._latest_train_collection = self._get_latest_collection(existing_train_collections, "train")
+        self._latest_val_collection = self._get_latest_collection(existing_val_collections, "val")
+
+    def _get_latest_collection(
+        self, collections: List[EntitiesCollectionInfo], expected_prefix: str
+    ) -> EntitiesCollectionInfo:
+        curr_collection = None
+        curr_idx = 0
+        for collection in collections:
+            parts = collection.name.split("_")
+            if len(parts) == 2:
+                prefix = parts[0].lower()
+                if prefix == expected_prefix:
+                    if parts[1].isdigit():
+                        collection_idx = int(parts[1])
+                        if collection_idx > curr_idx:
+                            curr_idx = collection_idx
+                            curr_collection = collection
+        return curr_collection
+
+
     def _detect_collections(self) -> bool:
         """Find collections with train and val prefixes and set them to train_val_splits"""
-        def _get_latest_collection(collections: List[EntitiesCollectionInfo]) -> EntitiesCollectionInfo:
-            curr_collection = None
-            curr_idx = 0
-            for collection in collections:
-                collection_idx = int(collection.name.rsplit('_', 1)[-1])
-                if collection_idx > curr_idx:
-                    curr_idx = collection_idx
-                    curr_collection = collection
-            return curr_collection
 
-        all_collections = self.api.entities_collection.get_list(self.project_id)
-        train_collections = []
-        val_collections = []
         collections_found = False
-        for collection in all_collections:
-            if collection.name.lower().startswith("train_"):
-                train_collections.append(collection)
-            elif collection.name.lower().startswith("val_"):
-                val_collections.append(collection)
-
-        train_collection = _get_latest_collection(train_collections)
-        val_collection = _get_latest_collection(val_collections)
-        if train_collection is not None and val_collection is not None:
-            self.train_val_splits.set_collections_splits([train_collection.id], [val_collection.id])
+        if self._latest_train_collection is not None and self._latest_val_collection is not None:
+            self.train_val_splits.set_collections_splits(
+                [self._latest_train_collection.id], [self._latest_val_collection.id]
+            )
             self.validator_text = Text("Train and val collections are detected", status="info")
             self.validator_text.show()
             collections_found = True
-            self._all_train_collections = train_collections
-            self._all_val_collections = val_collections
-            self._latest_train_collection = train_collection
-            self._latest_val_collection = val_collection
         else:
             self.validator_text = Text("")
             self.validator_text.hide()
             collections_found = False
         return collections_found
-        
+
     def _detect_datasets(self) -> bool:
         """Find datasets with train and val prefixes and set them to train_val_splits"""
+
         def _extend_with_nested(root_ds):
             nested = self.api.dataset.get_nested(self.project_id, root_ds.id)
             nested_ids = [ds.id for ds in nested]
@@ -407,7 +426,9 @@ class TrainValSplitsSelector:
         val_count = len(train_val_dataset_ids["val"])
 
         if train_count > 0 and val_count > 0:
-            self.train_val_splits.set_datasets_splits(train_val_dataset_ids["train"], train_val_dataset_ids["val"])
+            self.train_val_splits.set_datasets_splits(
+                train_val_dataset_ids["train"], train_val_dataset_ids["val"]
+            )
             datasets_found = True
 
         if train_count > 0 and val_count > 0:
@@ -415,7 +436,7 @@ class TrainValSplitsSelector:
                 message = "train and val datasets are detected"
             else:
                 message = "Multiple train and val datasets are detected. Check manually if selection is correct"
-            
+
             self.validator_text = Text(message, status="info")
             self.validator_text.show()
             datasets_found = True
@@ -423,4 +444,4 @@ class TrainValSplitsSelector:
             self.validator_text = Text("")
             self.validator_text.hide()
             datasets_found = False
-        return datasets_found
+        return datasets_found