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

supervisely/nn/live_training/artifacts_utils.py

@@ -0,0 +1,243 @@
+import os
+import shutil
+from pathlib import Path
+from typing import Optional, Dict
+from datetime import datetime
+import supervisely as sly
+from supervisely import logger
+from supervisely.template.live_training.live_training_generator import LiveTrainingGenerator
+import supervisely.io.json as sly_json
+from supervisely.nn.live_training.helpers import ClassMap
+import yaml
+
+
+def upload_artifacts(
+    api: sly.Api,
+    session_info: dict,
+    artifacts: dict,
+) -> str:
+    """
+    Upload artifacts to Team Files and generate experiment report.
+
+    Args:
+        session_info: Training session context
+            - team_id: Team ID
+            - task_id: Task ID
+            - project_id: Project ID
+            - framework_name: Framework name
+            - task_type: Task type
+            - class_map: Model class map
+            - start_time: Training start time string
+            - train_size: Final dataset size
+            - initial_samples: Number of initial samples
+            
+
+        artifacts: Framework-specific artifacts
+            - checkpoint_path: Path to checkpoint file
+            - checkpoint_info: Dict with {name, iteration, loss}
+            - config_path: Path to config file
+            - logs_dir: Path to TensorBoard logs or None
+            - model_config: Model configuration dict
+            - loss_history: Dict with loss history
+
+    Returns:
+        report_url: URL to experiment report
+    """
+    logger.info("Starting artifacts upload")
+
+    # Unpack session_info
+    team_id = session_info['team_id']
+    task_id = session_info['task_id']
+    project_id = session_info['project_id']
+    framework_name = session_info['framework_name']
+    task_type = session_info['task_type']
+    class_map: ClassMap = session_info['class_map']
+    model_meta = sly.ProjectMeta(obj_classes=class_map.obj_classes)
+    start_time = session_info['start_time']
+    train_size = session_info['train_size']
+    initial_samples = session_info.get('initial_samples', 0)
+    
+    # Unpack artifacts
+    checkpoint_path = artifacts['checkpoint_path']
+    checkpoint_info = artifacts['checkpoint_info']
+    config_path = artifacts['config_path']
+    logs_dir = artifacts.get('logs_dir')
+    model_config = artifacts['model_config']
+
+    work_dir = Path(os.path.dirname(checkpoint_path)).parent
+    output_dir = work_dir / "upload_artifacts"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    project_info = api.project.get_info_by_id(project_id)
+    project_name = project_info.name if project_info else "unknown"
+    model_name = f"Live training - {project_name}"
+    model_config['model_name'] = model_name
+    remote_dir = f"/experiments/live_training/{project_id}_{project_name}/{task_id}_{framework_name}/"
+    logger.info(f"Remote directory: {remote_dir}")
+
+    checkpoints_dir = output_dir / "checkpoints"
+    checkpoints_dir.mkdir(exist_ok=True)
+    checkpoint_dest = checkpoints_dir / checkpoint_info["name"]
+    shutil.copy2(checkpoint_path, checkpoint_dest)
+
+    state_json_src = checkpoint_path.replace('.pth', '_state.json')
+    if os.path.exists(state_json_src):
+        state_json_dest = str(checkpoint_dest).replace('.pth', '_state.json')
+        shutil.copy2(state_json_src, state_json_dest)
+
+    if config_path and os.path.exists(config_path):
+        config_dest = output_dir / os.path.basename(config_path)
+        shutil.copy2(config_path, config_dest)
+        model_files = {"config": config_dest.name}
+    else:
+        model_files = {}
+
+    sly_json.dump_json_file(
+        model_meta.to_json(),
+        str(output_dir / "model_meta.json")
+    )
+
+    hyperparams = {}
+    if config_path and os.path.exists(config_path):
+        try:
+            hyperparams = LiveTrainingGenerator.parse_hyperparameters(config_path)
+        except Exception as e:
+            logger.warning(f"Failed to parse hyperparameters: {e}")
+
+    with open(output_dir / "hyperparameters.yaml", 'w') as f:
+        yaml.dump(hyperparams, f, default_flow_style=False)
+
+    with open(output_dir / "open_app.lnk", 'w') as f:
+        f.write(f"/apps/sessions/{task_id}")
+
+    if logs_dir and os.path.exists(logs_dir):
+        logs_dest = output_dir / "logs"
+        if logs_dest.exists():
+            shutil.rmtree(logs_dest)
+        shutil.copytree(logs_dir, logs_dest)
+        logger.info(f"Logs copied from {logs_dir}")
+        has_logs = True
+    else:
+        logger.warning("No logs provided")
+        has_logs = False
+
+    logger.info("Uploading to Team Files")
+    api.file.upload_directory_fast(
+        team_id=team_id,
+        local_dir=str(output_dir),
+        remote_dir=remote_dir
+    )
+
+    experiment_info = {
+        "experiment_name": f"Live Training {task_type.capitalize()} - Task {task_id}",
+        "framework_name": framework_name,
+        "model_name": model_name,
+        "base_checkpoint": None,
+        "base_checkpoint_link": None,
+        "task_type": task_type,
+        "project_id": project_id,
+        "project_version": project_info.version if project_info else None,
+        "task_id": task_id,
+        "model_files": model_files,
+        "checkpoints": [f"checkpoints/{checkpoint_info['name']}"],
+        "best_checkpoint": checkpoint_info['name'],
+        "export": {},
+        "model_meta": "model_meta.json",
+        "hyperparameters": "hyperparameters.yaml",
+        "hyperparameters_id": None,
+        "artifacts_dir": remote_dir,
+        "datetime": start_time,
+        "experiment_report_id": None,
+        "evaluation_report_link": None,
+        "evaluation_metrics": {},
+        "primary_metric": None,
+        "logs": {"type": "tensorboard", "link": f"{remote_dir}logs/"} if has_logs else None,
+        "device": get_device_name(),
+        "training_duration": calculate_duration(start_time),
+        "train_collection_id": None,
+        "val_collection_id": None,
+        "project_preview": project_info.image_preview_url if project_info else None,
+        "train_size": train_size,
+        "initial_samples": initial_samples,
+        "val_size": 0,
+    }
+
+    checkpoints_info = [checkpoint_info]
+    loss_history = artifacts.get('loss_history', {})
+
+    session_info = {
+        "session_id": task_id,
+        "session_name": experiment_info["experiment_name"],
+        "project_id": project_id,
+        "start_time": experiment_info["datetime"],
+        "duration": experiment_info["training_duration"],
+        "artifacts_dir": remote_dir,
+        "logs_dir": f"{remote_dir}logs/" if has_logs else None,
+        "checkpoints": checkpoints_info,
+        "loss_history": loss_history,
+        "hyperparameters": hyperparams,
+        "status": "completed",
+        "device": experiment_info["device"],
+        "dataset_size": train_size,
+        "initial_samples": 0,
+        "samples_added": 0,
+        "final_size": train_size,
+        "train_size": train_size,
+        "initial_samples": initial_samples,
+        "val_size": 0,
+    }
+
+    report_dir = work_dir / "live_training_report"
+    report_dir.mkdir(exist_ok=True)
+
+    generator = LiveTrainingGenerator(
+        api=api,
+        session_info=session_info,
+        model_config=model_config,
+        model_meta=model_meta,
+        output_dir=str(report_dir),
+        team_id=team_id,
+        task_type=task_type,
+    )
+
+    generator.generate()
+    file_info = generator.upload_to_artifacts(os.path.join(remote_dir, "visualization"))
+
+    report_id = file_info if isinstance(file_info, int) else getattr(file_info, 'id', file_info)
+    report_url = f"{api.server_address}/nn/experiments/{report_id}"
+
+    logger.info(f"Report URL: {report_url}")
+
+    experiment_info["has_report"] = True
+    experiment_info["experiment_report_id"] = int(report_url.split('/')[-1])
+    response = api.task.set_output_experiment(task_id, experiment_info)
+
+    return report_url
+
+
+def calculate_duration(start_time: str) -> str:
+    """Calculate training duration in 'Xh Ym' format."""
+    try:
+        start_dt = datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S")
+        duration_sec = (datetime.now() - start_dt).total_seconds()
+        hours = int(duration_sec // 3600)
+        minutes = int((duration_sec % 3600) // 60)
+        return f"{hours}h {minutes}m"
+    except:
+        return "N/A"
+
+
+def get_device_name() -> str:
+    """Get GPU device name or 'cpu'."""
+    import torch # pylint: disable=import-error
+    if not os.path.exists("/dev/nvidia0"):
+        return "cpu"
+
+    try:
+        if torch.cuda.is_available():
+            device_id = int(os.getenv("CUDA_VISIBLE_DEVICES", "0").split(",")[0])
+            return torch.cuda.get_device_name(device_id)
+    except Exception as e:
+        logger.warning(f"Failed to get GPU name: {e}")
+
+    return "cuda"

supervisely/nn/live_training/checkpoint_utils.py

@@ -0,0 +1,229 @@
+import os
+import re
+from typing import Tuple, List, Optional
+import json
+from supervisely import logger
+from supervisely.nn.live_training.helpers import ClassMap
+from supervisely.io.json import load_json_file
+import supervisely as sly
+
+
+def validate_classes_exact_match(saved_classes: List[str], current_classes: List[str]):
+    """
+    Validate that saved and current classes match exactly.
+    Raises ValueError if they don't match.
+    """
+    if set(saved_classes) != set(current_classes):
+        raise ValueError(
+            f"Class names in checkpoint do not match current class names. "
+            f"Saved: {saved_classes}, Current: {current_classes}"
+        )
+
+
+def reorder_class_map(saved_classes: List[str], project_meta) -> ClassMap:
+    """
+    Create ClassMap with reordered classes matching checkpoint order.
+
+    Returns:
+        class_map: New ClassMap with correct order
+    """
+    class_mapping = {cls: idx for idx, cls in enumerate(saved_classes)}
+    logger.info(f"Class mapping: {class_mapping}")
+
+    # Create ClassMap from class names in checkpoint order
+    obj_classes = []
+    for name in saved_classes:
+        obj_class = project_meta.get_obj_class(name)
+        if obj_class is None:
+            raise ValueError(f"Class '{name}' not found in project metadata")
+        obj_classes.append(obj_class)
+
+    return ClassMap(obj_classes)
+
+
+def remove_classification_head(checkpoint_path: str) -> str:
+    """
+    Remove classification head weights from checkpoint and save modified version.
+    
+    Returns:
+        modified_path: Path to checkpoint without classification head
+    """
+    import torch  # pylint: disable=import-error
+    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+    state_dict = checkpoint.get('state_dict', {})
+    
+    keys_to_remove = []
+    for key in state_dict.keys():
+        if 'decode_head' in key or 'auxiliary_head' in key:
+            keys_to_remove.append(key)
+    
+    for key in keys_to_remove:
+        del state_dict[key]
+    
+    logger.info(f"Removed {len(keys_to_remove)} classification head parameters")
+    
+    modified_path = checkpoint_path.replace('.pth', '_headless.pth')
+    torch.save(checkpoint, modified_path)
+    
+    return modified_path
+
+
+def resolve_checkpoint(
+    checkpoint_mode: str,
+    selected_experiment_task_id: Optional[int],
+    class_map: ClassMap,
+    project_meta,
+    api,
+    team_id: int,
+    work_dir: str
+) -> Tuple[Optional[str], ClassMap, Optional[dict]]:
+    """
+    Main orchestrator function to resolve checkpoint loading based on mode.
+    
+    Args:
+        checkpoint_mode: One of 'scratch', 'finetune', 'continue'
+        selected_experiment_task_id: Task ID to load checkpoint from (required for finetune/continue)
+        class_map: Current ClassMap
+        project_meta: Project metadata
+        api: Supervisely API instance
+        team_id: Team ID
+        framework_name: Framework name (unused, kept for compatibility)
+        work_dir: Working directory for downloaded files
+    
+    Returns:
+        checkpoint_path: Path to checkpoint file (None for scratch mode)
+        class_map: Updated ClassMap (may be reordered for finetune/continue)
+        state: Training state dict (only for continue mode)
+    """
+    checkpoint_name = "latest.pth"
+    current_classes = [cls.name for cls in class_map.obj_classes]
+    logger.info(f"Checkpoint mode: {checkpoint_mode}")
+    
+    if checkpoint_mode == "scratch":
+        logger.info("Starting from pretrained weights (scratch mode)")
+        return None, class_map, None
+    
+    if selected_experiment_task_id is None:
+        raise ValueError(
+            f"selected_experiment_task_id must be provided when checkpoint_mode='{checkpoint_mode}'"
+        )
+    
+    # Get experiment info
+    task_info = api.task.get_info_by_id(selected_experiment_task_id)
+    experiment_info = task_info["meta"]["output"]["experiment"]["data"]
+
+    artifacts_dir = experiment_info["artifacts_dir"]
+    model_meta_filename = experiment_info.get("model_meta", "model_meta.json")
+
+    # Setup local paths
+    local_dir = os.path.join(work_dir, 'downloaded_checkpoints')
+    os.makedirs(local_dir, exist_ok=True)
+
+    # Download checkpoint
+    remote_checkpoint = f"{artifacts_dir}checkpoints/{checkpoint_name}"
+    local_checkpoint = os.path.join(local_dir, checkpoint_name)
+    
+    logger.info(f"Downloading checkpoint from {remote_checkpoint}")
+    api.file.download(team_id, remote_checkpoint, local_checkpoint)
+    logger.info(f"Checkpoint downloaded to {local_checkpoint}")
+    
+    # Download model_meta.json
+    remote_model_meta = f"{artifacts_dir}{model_meta_filename}"
+    local_model_meta = os.path.join(local_dir, 'model_meta.json')
+    
+    logger.info(f"Downloading model_meta from {remote_model_meta}")
+    api.file.download(team_id, remote_model_meta, local_model_meta)
+    
+    # Load saved classes
+    model_meta_json = load_json_file(local_model_meta)
+    saved_project_meta = sly.ProjectMeta.from_json(model_meta_json)
+    saved_classes = [cls.name for cls in saved_project_meta.obj_classes]
+    
+    logger.info(f"Saved classes: {saved_classes}")
+    logger.info(f"Current classes: {current_classes}")
+    
+    # Finetune mode - flexible class handling
+    if checkpoint_mode == "finetune":
+        saved_set = set(saved_classes)
+        current_set = set(current_classes)
+        
+        if saved_set == current_set:
+            if saved_classes != current_classes:
+                logger.info("Class order differs. Reordering classes")
+                class_map = reorder_class_map(saved_classes, project_meta)
+            else:
+                logger.info("Class names match exactly")
+            return local_checkpoint, class_map, None
+        
+        elif len(saved_classes) == len(current_classes):
+            # logger.info("Class names differ but count matches. Removing classification head")
+            # modified_checkpoint = remove_classification_head(local_checkpoint)
+            logger.warning("Class names differ but count matches. Classification head will be kept as is")
+            return local_checkpoint, class_map, None
+            
+        else:
+            logger.info("Classes differ completely. Starting from checkpoint")
+            return local_checkpoint, class_map, None
+    
+    # Continue mode - strict matching required
+    elif checkpoint_mode == "continue":
+        logger.info(f"Continue mode: loading from task_id={selected_experiment_task_id}")
+        validate_classes_exact_match(saved_classes, current_classes)
+
+        # Download state JSON
+        state_filename = checkpoint_name.replace('.pth', '_state.json')
+        remote_state = f"{artifacts_dir}checkpoints/{state_filename}"
+        local_state = local_checkpoint.replace('.pth', '_state.json')
+        
+        logger.info(f"Downloading state from {remote_state}")
+        api.file.download(team_id, remote_state, local_state)
+        
+        # Use existing utility function
+        state = load_state_json(local_checkpoint)
+        
+        logger.info(f"State loaded with {state.get('dataset_size', 0)} samples at iter {state.get('iter', 0)}")
+        return local_checkpoint, class_map, state
+    
+    else:
+        raise ValueError(
+            f"Invalid checkpoint_mode='{checkpoint_mode}'. Valid: 'scratch', 'finetune', 'continue'"
+        )
+
+def save_state_json(state: dict, checkpoint_path: str):
+    """
+    Save training state as JSON file next to checkpoint.
+    
+    Args:
+        state: State dict from LiveTraining.state()
+        checkpoint_path: Path to .pth checkpoint file
+    """
+    state_path = checkpoint_path.replace('.pth', '_state.json')
+    
+    with open(state_path, 'w') as f:
+        json.dump(state, f, indent=2)
+    
+    logger.info(f"State saved to {state_path}")
+
+def load_state_json(checkpoint_path: str) -> dict:
+    """
+    Load training state from JSON file next to checkpoint.
+    
+    Args:
+        checkpoint_path: Path to .pth checkpoint file
+        
+    Returns:
+        state: State dict for LiveTraining.load_state()
+    """
+    state_path = checkpoint_path.replace('.pth', '_state.json')
+    
+    if not os.path.exists(state_path):
+        raise ValueError(
+            f"State file not found: {state_path}. "
+            f"This checkpoint may not support 'continue' mode."
+        )
+    
+    with open(state_path, 'r') as f:
+        state = json.load(f)
+    
+    logger.info(f"State loaded from {state_path}")
+    return state

supervisely/nn/live_training/dynamic_sampler.py

@@ -0,0 +1,44 @@
+import random
+import time
+from collections.abc import Sized
+
+
+class DynamicSampler:
+    """
+    A sampler that dynamically adjusts to the size of a dataset that grows over time.
+    Implements torch.utils.data.Sampler interface.
+    """
+    
+    def __init__(self, dataset: Sized, shuffle: bool = True, seed: int = 0):
+        self.dataset = dataset
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def __iter__(self):
+        remaining_indices = []
+        last_known_len = 0
+        
+        # Wait for first samples
+        while len(self.dataset) == 0:
+            time.sleep(0.1)
+        
+        while True:
+            current_len = len(self.dataset)
+            
+            if current_len > last_known_len:
+                new_indices = list(range(last_known_len, current_len))
+                if self.shuffle:
+                    random.shuffle(new_indices)
+                remaining_indices.extend(new_indices)
+                last_known_len = current_len
+            
+            if not remaining_indices:
+                # Reshuffle existing data
+                remaining_indices = list(range(current_len))
+                if self.shuffle:
+                    random.shuffle(remaining_indices)
+            
+            yield remaining_indices.pop()
+    
+    def __len__(self):
+        return len(self.dataset)

supervisely/nn/live_training/helpers.py

@@ -0,0 +1,14 @@
+from typing import List, Union
+import supervisely as sly
+
+
+class ClassMap:
+    def __init__(self, obj_classes: Union[sly.ObjClassCollection, List[sly.ObjClass]]):
+        self.obj_classes = obj_classes
+        self.class2idx = {obj_class.name: idx for idx, obj_class in enumerate(self.obj_classes)}
+        self.idx2class = {idx: obj_class.name for idx, obj_class in enumerate(self.obj_classes)}
+        self.classes = [obj_class.name for obj_class in self.obj_classes]
+        self.sly_ids = [obj_class.sly_id for obj_class in self.obj_classes]
+    
+    def __len__(self):
+        return len(self.obj_classes)

supervisely/nn/live_training/incremental_dataset.py

@@ -0,0 +1,146 @@
+from typing import Dict, Any
+from pathlib import Path
+from PIL import Image
+import numpy as np
+import supervisely as sly
+import cv2
+
+
+class IncrementalDataset:
+    """
+    1. Save images on disk
+    2. Store annotations in SLY/COCO format. Handle case for Segmentation task
+    3. Implement indexing, adding, and updating samples
+    4. __getitem__ to retrieve samples by index
+    """
+    def __init__(
+            self,
+            class2idx: dict,
+            data_dir: str,
+            save_masks_as_images: bool = False,
+        ):
+        self.class2idx = class2idx
+        self.data_dir = Path(data_dir)
+        self.save_masks_as_images = save_masks_as_images
+        self.images_dir = self.data_dir / "images"
+        self.images_dir.mkdir(parents=True, exist_ok=True)
+        if self.save_masks_as_images:
+            self.masks_dir = self.data_dir / "masks"
+            self.masks_dir.mkdir(parents=True, exist_ok=True)
+        self.samples: Dict[int, dict] = {}
+        self.samples_list = []
+
+    def add(
+            self,
+            image_id: int,
+            image_np: np.ndarray,
+            annotation: sly.Annotation,
+            image_name: str
+        ) -> dict:
+        if image_id in self.samples:
+            raise ValueError(f"Cannot add sample: Image ID {image_id} already exists in the dataset.")
+        image_name = f"{image_id} {image_name}"
+        w, h = image_np.shape[1], image_np.shape[0]
+        img_size = (w, h)
+        image_path = self._save_img(image_np, image_name)
+        mask_path = None
+        if self.save_masks_as_images:
+            mask_path = self._save_mask(annotation, image_name)
+        sample = self._format_sample(
+            image_id,
+            annotation,
+            img_size,
+            image_path,
+            mask_path
+        )
+        assert isinstance(sample, dict), "Sample must be a dict."
+        # add extra fields for internal use
+        sample['image_path'] = image_path
+        sample['size'] = img_size
+        if mask_path is not None:
+            sample['mask_path'] = mask_path
+        # add to dataset
+        self.samples[image_id] = sample
+        self.samples_list.append(sample)
+        return sample
+    
+    def update(
+            self,
+            image_id: int,
+            annotation: sly.Annotation,
+        ) -> dict:
+        if image_id not in self.samples:
+            raise ValueError(f"Cannot update sample: Image ID {image_id} does not exist in the dataset.")
+        sample = self.samples[image_id]
+        new_sample = self._format_sample(
+            image_id,
+            annotation,
+            sample['size'],
+            sample['image_path'],
+            sample.get('mask_path')
+        )
+        sample.update(new_sample)
+        return sample
+    
+    def add_or_update(
+            self,
+            image_id: int,
+            image_np: np.ndarray,
+            annotation: sly.Annotation,
+            image_name: str
+        ) -> dict:
+        if image_id not in self.samples:
+            return self.add(image_id, image_np, annotation, image_name)
+        else:
+            return self.update(image_id, annotation)
+    
+    def _format_sample(
+            self,
+            image_id: int,
+            annotation: sly.Annotation, 
+            image_size: tuple,
+            image_path: str,
+            mask_path: str = None
+        ) -> dict:
+        sample = {
+            'image_id': image_id,
+            'width': image_size[0],
+            'height': image_size[1],
+            'annotations': annotation.to_coco(annotation.image_id, self.class2idx)[0],
+            'image_path': image_path,
+            'mask_path': mask_path
+        }
+        return sample
+    
+    def _save_img(self, image_np: np.ndarray, image_name: str) -> str:
+        image = Image.fromarray(image_np).convert('RGB')
+        image_path = str(self.images_dir / image_name)
+        image.save(image_path)
+        return image_path
+
+    def _save_mask(self, annotation: sly.Annotation, image_name: str) -> str:
+
+        mapping = {label.obj_class: label.obj_class for label in annotation.labels}
+        ann_nonoverlap = annotation.to_nonoverlapping_masks(mapping)
+        h, w = annotation.img_size
+        mask = np.zeros((h, w), dtype=np.uint8)
+        
+        for label in ann_nonoverlap.labels:
+            class_name = label.obj_class.name
+            class_id = self.class2idx.get(class_name)
+            if class_id is not None:
+                label.geometry.draw(mask, color=class_id)
+        
+        mask_name = Path(image_name).stem + '.png' 
+        mask_path = str(self.masks_dir / mask_name)
+        cv2.imwrite(mask_path, mask)
+        
+        return mask_path
+
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def get_image_ids(self) -> list:
+        """Get list of image IDs in dataset"""
+        return list(self.samples.keys())
+