synapse-sdk 1.0.0a11__py3-none-any.whl → 2026.1.1b2__py3-none-any.whl

synapse_sdk/plugins/actions/dataset/action.py

@@ -0,0 +1,471 @@
+"""Dataset action with download and convert operations.
+
+A single action class that handles both dataset download and format conversion,
+selected via the operation parameter. Designed for pipeline composition.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from enum import StrEnum
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from synapse_sdk.plugins.action import BaseAction
+from synapse_sdk.plugins.datasets.converters import DatasetFormat, get_converter
+from synapse_sdk.plugins.datasets.formats.dm import DMVersion
+from synapse_sdk.plugins.enums import PluginCategory
+from synapse_sdk.plugins.types import YOLODataset
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.backend import BackendClient
+
+
+class DatasetOperation(StrEnum):
+    """Dataset operation types."""
+
+    DOWNLOAD = 'download'
+    CONVERT = 'convert'
+
+
+class DatasetParams(BaseModel):
+    """Parameters for DatasetAction.
+
+    The operation field determines which operation to perform:
+    - download: Downloads dataset from backend
+    - convert: Converts dataset from one format to another
+
+    Attributes:
+        operation: Which operation to perform.
+        dataset_id: Data collection ID (for download).
+        splits: Split definitions for categorized download.
+        path: Source dataset path (for convert, or set by download).
+        source_format: Source format (for convert).
+        target_format: Target format (for convert).
+        dm_version: Datamaker version (for convert from DM).
+        output_dir: Output directory (optional for both).
+        is_categorized: Whether dataset has train/valid/test splits.
+    """
+
+    operation: DatasetOperation = DatasetOperation.DOWNLOAD
+
+    # Download params
+    dataset_id: int | None = Field(default=None, description='Data collection ID')
+    splits: dict[str, dict[str, Any]] | None = Field(
+        default=None,
+        description='Split definitions: {"train": {...filters}, "valid": {...}}',
+    )
+
+    # Convert params
+    path: Path | str | None = Field(default=None, description='Dataset path')
+    source_format: str = Field(default='dm_v2', description='Source format')
+    target_format: str = Field(default='yolo', description='Target format')
+    dm_version: str = Field(default='v2', description='Datamaker version')
+
+    # Shared params
+    output_dir: Path | str | None = Field(default=None, description='Output directory')
+    is_categorized: bool = Field(default=False, description='Has splits')
+
+
+class DatasetResult(BaseModel):
+    """Result from DatasetAction.
+
+    Contains paths and metadata about the processed dataset.
+
+    Attributes:
+        path: Path to dataset directory.
+        format: Dataset format (e.g., 'dm_v2', 'yolo').
+        is_categorized: Whether dataset has splits.
+        config_path: Path to config file (e.g., dataset.yaml for YOLO).
+        count: Number of items processed.
+        source_path: Original source path (for convert).
+        data_path: Computed property returning config_path if set, otherwise path.
+    """
+
+    path: Path
+    format: str
+    is_categorized: bool = False
+    config_path: Path | None = None
+    count: int | None = None
+    source_path: Path | None = None
+
+    @property
+    def data_path(self) -> Path:
+        """Returns config_path if set, otherwise path.
+
+        Use this for downstream actions that need a single path
+        to the dataset (e.g., training with YOLO format).
+        """
+        return self.config_path if self.config_path is not None else self.path
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class DatasetAction(BaseAction[DatasetParams]):
+    """Dataset action with download and convert operations.
+
+    A unified action for dataset operations that can be used in pipelines.
+    The operation is determined by the params.operation field.
+
+    Type declarations:
+        - input_type: None (accepts initial params)
+        - output_type: Dynamic based on operation and target_format
+          - download: 'dm_dataset'
+          - convert to yolo: 'yolo_dataset'
+          - convert to coco: 'coco_dataset'
+
+    For download:
+        - Requires: dataset_id
+        - Optional: splits, output_dir
+        - Returns: path, format='dm_v2', is_categorized, count
+
+    For convert:
+        - Requires: path, target_format
+        - Optional: source_format, dm_version, output_dir
+        - Returns: path, format, config_path, source_path
+
+    Example:
+        >>> # Standalone usage
+        >>> action = DatasetAction(
+        ...     DatasetParams(operation='download', dataset_id=123),
+        ...     ctx,
+        ... )
+        >>> result = action.execute()
+        >>>
+        >>> # Pipeline usage
+        >>> pipeline = ActionPipeline([DatasetAction, DatasetAction, TrainAction])
+        >>> result = pipeline.execute({
+        ...     'operation': 'download',
+        ...     'dataset_id': 123,
+        ...     'target_format': 'yolo',  # Used by second DatasetAction
+        ... }, ctx)
+    """
+
+    category = PluginCategory.NEURAL_NET
+
+    # Input type is flexible (accepts various initial params)
+    input_type = None
+    # Output type: use YOLODataset for convert (most common), DMv2Dataset for download
+    # For precise typing, use separate DownloadAction/ConvertAction classes
+    output_type = YOLODataset  # Default assumes convert to YOLO
+
+    result_model = DatasetResult
+
+    @property
+    def client(self) -> BackendClient:
+        """Backend client from context."""
+        if self.ctx.client is None:
+            raise RuntimeError('No backend client in context')
+        return self.ctx.client
+
+    def execute(self) -> DatasetResult:
+        """Execute the dataset operation based on params.operation."""
+        if self.params.operation == DatasetOperation.DOWNLOAD:
+            return self.download()
+        elif self.params.operation == DatasetOperation.CONVERT:
+            return self.convert()
+        else:
+            raise ValueError(f'Unknown operation: {self.params.operation}')
+
+    def download(self) -> DatasetResult:
+        """Download dataset from backend.
+
+        Downloads data units from a data collection and saves them
+        locally in Datamaker format (json/ + original_files/).
+
+        Returns:
+            DatasetResult with path, format, count.
+
+        Raises:
+            ValueError: If dataset_id not provided.
+        """
+        from synapse_sdk.utils.file import get_temp_path
+
+        if self.params.dataset_id is None:
+            raise ValueError('dataset_id is required for download operation')
+
+        dataset_id = self.params.dataset_id
+        splits = self.params.splits
+        is_categorized = splits is not None and len(splits) > 0
+
+        # Determine output directory
+        output_dir = Path(self.params.output_dir) if self.params.output_dir else get_temp_path(f'datasets/{dataset_id}')
+        output_dir = Path(output_dir)
+
+        # Report initial progress
+        self.set_progress(0, 100, 'init')
+
+        # Get collection info
+        collection = self.client.get_data_collection(dataset_id)
+        self.log(
+            'download_start',
+            {
+                'dataset_id': dataset_id,
+                'collection_name': collection.get('name'),
+                'is_categorized': is_categorized,
+            },
+        )
+
+        # Report collection fetched
+        self.set_progress(1, 100, 'init')
+
+        total_downloaded = 0
+
+        if is_categorized and splits:
+            # Download each split separately
+            for split_name, filters in splits.items():
+                split_dir = output_dir / split_name
+                count = self._download_split(
+                    dataset_id=dataset_id,
+                    output_dir=split_dir,
+                    filters=filters or {},
+                )
+                total_downloaded += count
+                self.log(
+                    'split_downloaded',
+                    {
+                        'split': split_name,
+                        'count': count,
+                    },
+                )
+        else:
+            # Download all data units
+            total_downloaded = self._download_split(
+                dataset_id=dataset_id,
+                output_dir=output_dir,
+                filters={},
+            )
+
+        self.log(
+            'download_complete',
+            {
+                'path': str(output_dir),
+                'total_units': total_downloaded,
+            },
+        )
+
+        return DatasetResult(
+            path=output_dir,
+            format='dm_v2',
+            is_categorized=is_categorized,
+            count=total_downloaded,
+        )
+
+    def _download_split(
+        self,
+        dataset_id: int,
+        output_dir: Path,
+        filters: dict[str, Any],
+        max_workers: int = 10,
+    ) -> int:
+        """Download a single split of the dataset."""
+        # Create output directories
+        json_dir = output_dir / 'json'
+        files_dir = output_dir / 'original_files'
+        json_dir.mkdir(parents=True, exist_ok=True)
+        files_dir.mkdir(parents=True, exist_ok=True)
+
+        # Report fetching data units
+        self.set_progress(2, 100, 'fetch')
+
+        # List data units
+        params = {'data_collection': dataset_id, **filters}
+        data_units_gen, total_count = self.client.list_data_units(
+            params=params,
+            list_all=True,
+        )
+
+        # Report data units fetched
+        self.set_progress(5, 100, 'fetch')
+        self.log('data_units_listed', {'total_count': total_count})
+
+        downloaded = 0
+
+        def download_unit(unit: dict) -> bool:
+            """Download a single data unit."""
+            try:
+                unit_id = unit.get('id') or unit.get('data_unit_id')
+                files = unit.get('files', {})
+
+                # Build DM v2 JSON structure
+                dm_json = self._build_dm_json(unit)
+
+                # Determine base name from first file or unit ID
+                base_name = None
+                for file_info in files.values():
+                    if isinstance(file_info, dict):
+                        file_path = file_info.get('path') or file_info.get('url', '')
+                    else:
+                        file_path = str(file_info)
+                    if file_path:
+                        base_name = Path(file_path).stem
+                        break
+
+                if not base_name:
+                    base_name = str(unit_id)
+
+                # Save JSON
+                json_path = json_dir / f'{base_name}.json'
+                json_path.write_text(json.dumps(dm_json, indent=2, ensure_ascii=False))
+
+                # Copy/download files
+                for file_name, file_info in files.items():
+                    if isinstance(file_info, dict):
+                        file_path = file_info.get('path')
+                    else:
+                        file_path = str(file_info)
+
+                    if file_path and Path(file_path).exists():
+                        dest = files_dir / Path(file_path).name
+                        if not dest.exists():
+                            shutil.copy(file_path, dest)
+
+                return True
+            except Exception as e:
+                self.log('download_unit_error', {'unit_id': unit_id, 'error': str(e)})
+                return False
+
+        # Process units with thread pool
+        # Note: data_units_gen is a lazy generator that fetches pages from API
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = []
+            units_fetched = 0
+
+            # Submit downloads as we iterate (reports fetch progress)
+            for unit in data_units_gen:
+                futures.append(executor.submit(download_unit, unit))
+                units_fetched += 1
+                # Report fetch progress (5-50% range)
+                fetch_progress = 5 + (units_fetched / total_count) * 45
+                if units_fetched % 10 == 0 or units_fetched == total_count:
+                    self.set_progress(int(fetch_progress), 100, 'fetch')
+
+            # Process completed downloads (50-100% range)
+            for i, future in enumerate(as_completed(futures)):
+                if future.result():
+                    downloaded += 1
+                # Report download progress (50-100% range)
+                download_progress = 50 + ((i + 1) / total_count) * 50
+                self.set_progress(int(download_progress), 100, 'download')
+
+        return downloaded
+
+    def _build_dm_json(self, unit: dict) -> dict[str, Any]:
+        """Build Datamaker v2 JSON from a data unit."""
+        annotations = unit.get('annotations', {})
+
+        dm_image: dict[str, list] = {
+            'bounding_box': [],
+            'polygon': [],
+            'polyline': [],
+            'keypoint': [],
+            'relation': [],
+            'group': [],
+        }
+
+        if isinstance(annotations, dict):
+            for key in dm_image.keys():
+                if key in annotations:
+                    dm_image[key] = annotations[key]
+
+        # Build classification map from annotations
+        classifications: dict[str, set[str]] = {}
+        for ann_type, anns in dm_image.items():
+            if anns:
+                classifications[ann_type] = set()
+                for ann in anns:
+                    if 'classification' in ann:
+                        classifications[ann_type].add(ann['classification'])
+
+        return {
+            'classification': {k: sorted(v) for k, v in classifications.items() if v},
+            'images': [dm_image],
+        }
+
+    def convert(self) -> DatasetResult:
+        """Convert dataset from one format to another.
+
+        Converts the dataset at params.path to params.target_format.
+
+        Returns:
+            DatasetResult with converted path, format, config_path.
+
+        Raises:
+            ValueError: If path not provided.
+        """
+        if self.params.path is None:
+            raise ValueError('path is required for convert operation')
+
+        source_path = Path(self.params.path)
+        if not source_path.exists():
+            raise FileNotFoundError(f'Dataset path does not exist: {source_path}')
+
+        # Parse formats
+        target_format = DatasetFormat(self.params.target_format)
+        dm_version = DMVersion.V1 if self.params.dm_version == 'v1' else DMVersion.V2
+
+        # Determine source format
+        source_format_str = self.params.source_format
+        if source_format_str in ('dm_v1', 'dm_v2', 'dm'):
+            src_format = DatasetFormat.DM_V1 if dm_version == DMVersion.V1 else DatasetFormat.DM_V2
+        else:
+            src_format = DatasetFormat(source_format_str)
+
+        # Determine output directory
+        if self.params.output_dir:
+            output_dir = Path(self.params.output_dir)
+        else:
+            output_dir = source_path.parent / f'{source_path.name}_{target_format.value}'
+
+        self.log(
+            'convert_start',
+            {
+                'source_path': str(source_path),
+                'source_format': src_format.value,
+                'target_format': target_format.value,
+                'is_categorized': self.params.is_categorized,
+            },
+        )
+
+        # Get converter and run conversion
+        converter = get_converter(
+            source=src_format,
+            target=target_format,
+            root_dir=source_path,
+            is_categorized=self.params.is_categorized,
+            dm_version=dm_version,
+        )
+
+        converter.convert()
+        converter.save_to_folder(output_dir)
+
+        # Determine config path
+        config_path = None
+        if target_format == DatasetFormat.YOLO:
+            config_path = output_dir / 'dataset.yaml'
+            if not config_path.exists():
+                config_path = None
+
+        self.log(
+            'convert_complete',
+            {
+                'output_path': str(output_dir),
+                'config_path': str(config_path) if config_path else None,
+            },
+        )
+
+        return DatasetResult(
+            path=output_dir,
+            format=target_format.value,
+            is_categorized=self.params.is_categorized,
+            config_path=config_path,
+            source_path=source_path,
+        )
+
+
+__all__ = ['DatasetAction', 'DatasetOperation', 'DatasetParams', 'DatasetResult']

synapse_sdk/plugins/actions/export/__init__.py

@@ -0,0 +1,55 @@
+"""Export action module with optional workflow step support.
+
+Provides the export action base class:
+    - BaseExportAction: Base class for export workflows
+    - ExportContext: Export-specific context extending BaseStepContext
+    - ExportProgressCategories: Standard progress category names
+
+For step infrastructure (BaseStep, StepRegistry, Orchestrator),
+use the steps module:
+    from synapse_sdk.plugins.steps import BaseStep, StepRegistry
+
+Example (simple execute):
+    >>> class MyExportAction(BaseExportAction[MyParams]):
+    ...     def get_filtered_results(self, filters: dict) -> tuple[Any, int]:
+    ...         return self.client.get_assignments(filters)
+    ...
+    ...     def execute(self) -> dict[str, Any]:
+    ...         results, count = self.get_filtered_results(self.params.filter)
+    ...         # ... export items ...
+    ...         return {'exported': count}
+
+Example (step-based):
+    >>> from synapse_sdk.plugins.steps import BaseStep, StepResult
+    >>>
+    >>> class FetchResultsStep(BaseStep[ExportContext]):
+    ...     @property
+    ...     def name(self) -> str:
+    ...         return 'fetch_results'
+    ...
+    ...     @property
+    ...     def progress_weight(self) -> float:
+    ...         return 0.2
+    ...
+    ...     def execute(self, context: ExportContext) -> StepResult:
+    ...         context.results, context.total_count = fetch_data(context.params)
+    ...         return StepResult(success=True)
+    >>>
+    >>> class MyExportAction(BaseExportAction[MyParams]):
+    ...     def setup_steps(self, registry) -> None:
+    ...         registry.register(FetchResultsStep())
+    ...         registry.register(ProcessStep())
+    ...         registry.register(FinalizeStep())
+"""
+
+from synapse_sdk.plugins.actions.export.action import (
+    BaseExportAction,
+    ExportProgressCategories,
+)
+from synapse_sdk.plugins.actions.export.context import ExportContext
+
+__all__ = [
+    'BaseExportAction',
+    'ExportContext',
+    'ExportProgressCategories',
+]