Create training dataset (#4998)

Author: itallix

application/backend/app/api/dependencies.py

@@ -217,15 +217,17 @@ def get_project_service(
 ) -> ProjectService:
     """Provides a ProjectService instance for managing projects."""
     return ProjectService(
-        data_dir=data_dir, db_session=db, label_service=label_service, pipeline_service=pipeline_service
+        data_dir=data_dir, label_service=label_service, pipeline_service=pipeline_service, db_session=db
     )
 
 
 def get_dataset_service(
-    data_dir: Annotated[Path, Depends(get_data_dir)], db: Annotated[Session, Depends(get_db)]
+    data_dir: Annotated[Path, Depends(get_data_dir)],
+    label_service: Annotated[LabelService, Depends(get_label_service)],
+    db: Annotated[Session, Depends(get_db)],
 ) -> DatasetService:
     """Provides a DatasetService instance."""
-    return DatasetService(data_dir=data_dir, db_session=db)
+    return DatasetService(data_dir=data_dir, label_service=label_service, db_session=db)
 
 
 def get_project(

application/backend/app/api/routers/datasets.py

@@ -20,7 +20,12 @@
 from app.models import DatasetItemAnnotationStatus, DatasetItemSubset
 from app.schemas import ProjectView
 from app.services import DatasetService, ResourceNotFoundError
-from app.services.dataset_service import AnnotationValidationError, InvalidImageError, SubsetAlreadyAssignedError
+from app.services.dataset_service import (
+    AnnotationValidationError,
+    DatasetItemFilters,
+    InvalidImageError,
+    SubsetAlreadyAssignedError,
+)
 
 router = APIRouter(prefix="/api/projects/{project_id}/dataset/items", tags=["Datasets"])
 
@@ -122,14 +127,16 @@ def list_dataset_items(  # noqa: PLR0913
         subset=subset,
     )
     dataset_items = dataset_service.list_dataset_items(
-        project=project,
-        limit=limit,
-        offset=offset,
-        start_date=start_date,
-        end_date=end_date,
-        annotation_status=annotation_status,
-        label_ids=labels,
-        subset=subset,
+        project_id=project.id,
+        filters=DatasetItemFilters(
+            limit=limit,
+            offset=offset,
+            start_date=start_date,
+            end_date=end_date,
+            annotation_status=annotation_status,
+            label_ids=labels,
+            subset=subset,
+        ),
     )
     return DatasetItemsWithPagination(
         items=[DatasetItemView.model_validate(dataset_item, from_attributes=True) for dataset_item in dataset_items],
@@ -157,7 +164,7 @@ def get_dataset_item(
 ) -> DatasetItemView:
     """Get information about a specific dataset item"""
     try:
-        dataset_item = dataset_service.get_dataset_item_by_id(project=project, dataset_item_id=dataset_item_id)
+        dataset_item = dataset_service.get_dataset_item_by_id(project_id=project.id, dataset_item_id=dataset_item_id)
         return DatasetItemView.model_validate(dataset_item, from_attributes=True)
     except ResourceNotFoundError as e:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
@@ -179,7 +186,7 @@ def get_dataset_item_binary(
     """Get dataset item binary content"""
     try:
         binary_path = dataset_service.get_dataset_item_binary_path_by_id(
-            project=project, dataset_item_id=dataset_item_id
+            project_id=project.id, dataset_item_id=dataset_item_id
         )
         return FileResponse(path=binary_path)
     except ResourceNotFoundError as e:
@@ -281,7 +288,7 @@ def get_dataset_item_annotations(
 ) -> DatasetItemAnnotations:
     """Get the dataset item annotations"""
     try:
-        dataset_item = dataset_service.get_dataset_item_by_id(project=project, dataset_item_id=dataset_item_id)
+        dataset_item = dataset_service.get_dataset_item_by_id(project_id=project.id, dataset_item_id=dataset_item_id)
         if dataset_item.annotation_data is None:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND, detail="Dataset item has not been annotated yet."
@@ -336,7 +343,7 @@ def assign_dataset_item_subset(
     """Assign dataset item subset"""
     try:
         dataset_item = dataset_service.assign_dataset_item_subset(
-            project=project, dataset_item_id=dataset_item_id, subset=subset_config.subset
+            project_id=project.id, dataset_item_id=dataset_item_id, subset=subset_config.subset
         )
         return DatasetItemView.model_validate(dataset_item, from_attributes=True)
     except ResourceNotFoundError as e:

application/backend/app/api/routers/jobs.py

@@ -65,7 +65,7 @@ async def submit_job(
                     params=TrainingParams(
                         model_architecture_id=job_request.parameters.model_architecture_id,
                         parent_model_revision_id=job_request.parameters.parent_model_revision_id,
-                        task_type=project.task.task_type,
+                        task=project.task,
                     ),
                 )
             case _:

application/backend/app/core/logging/__init__.py

@@ -2,10 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .config import LogConfig
+from .handlers import InterceptHandler
 from .setup import setup_logging, setup_uvicorn_logging
 from .utils import logging_ctx
 
 __all__ = [
+    "InterceptHandler",
     "LogConfig",
     "logging_ctx",
     "setup_logging",

application/backend/app/lifecycle.py

@@ -19,6 +19,7 @@
 from app.db import MigrationManager, get_db_session
 from app.scheduler import Scheduler
 from app.schemas.job import JobType
+from app.services import DatasetService, LabelService
 from app.services.base_weights_service import BaseWeightsService
 from app.services.data_collect import DataCollector
 from app.services.event.event_bus import EventBus
@@ -48,13 +49,16 @@ def setup_job_controller(data_dir: Path, max_parallel_jobs: int) -> tuple[JobQue
     base_weights_service = BaseWeightsService(data_dir=data_dir)
     subset_service = SubsetService()
     subset_assigner = SubsetAssigner()
+    label_service = LabelService()
+    dataset_service = DatasetService(data_dir=data_dir, label_service=label_service)
     job_runnable_factory.register(
         JobType.TRAIN,
         partial(
             OTXTrainer,
             base_weights_service=base_weights_service,
             subset_service=subset_service,
             subset_assigner=subset_assigner,
+            dataset_service=dataset_service,
             data_dir=data_dir,
             db_session_factory=get_db_session,
         ),

application/backend/app/main.py

@@ -8,6 +8,7 @@
 #  - docker compose up
 
 import importlib
+import logging
 import pkgutil
 from pathlib import Path
 
@@ -18,10 +19,12 @@
 from loguru import logger
 
 from app.api import routers
+from app.core.logging import InterceptHandler
 from app.lifecycle import lifespan
 from app.settings import get_settings
 
 settings = get_settings()
+logging.basicConfig(handlers=[InterceptHandler()], level=settings.log_level, force=True)
 app = FastAPI(
     title=settings.app_name,
     version=settings.version,

application/backend/app/services/__init__.py

@@ -3,12 +3,14 @@
 
 from .active_model_service import ActiveModelService
 from .base import (
+    BaseSessionManagedService,
     ResourceInUseError,
     ResourceNotFoundError,
     ResourceType,
     ResourceWithIdAlreadyExistsError,
     ResourceWithNameAlreadyExistsError,
 )
+from .base_weights_service import BaseWeightsService
 from .dataset_service import DatasetService
 from .dispatch_service import DispatchService
 from .label_service import LabelService
@@ -24,6 +26,8 @@
 
 __all__ = [
     "ActiveModelService",
+    "BaseSessionManagedService",
+    "BaseWeightsService",
     "DatasetService",
     "DispatchService",
     "LabelService",

application/backend/app/services/base.py

@@ -1,8 +1,12 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from abc import ABC
+from collections.abc import Callable
 from enum import StrEnum
 
+from sqlalchemy.orm import Session
+
 
 class ResourceType(StrEnum):
     """Enumeration for resource types."""
@@ -55,3 +59,65 @@ class ResourceWithIdAlreadyExistsError(ResourceError):
     def __init__(self, resource_type: ResourceType, resource_id: str, message: str | None = None):
         msg = message or f"{resource_type} with ID '{resource_id}' already exists."
         super().__init__(resource_type, resource_id, msg)
+
+
+class BaseSessionManagedService(ABC):
+    """
+    Base class for services that require a managed database session.
+
+    This class supports deferred database session initialization, allowing services
+    to be instantiated without an immediate database connection. The session can be
+    provided either at construction time or injected later via `set_db_session()`.
+
+    This pattern is useful in scenarios where:
+    - Services need to be created before database context is available
+    - Database session management is handled externally (e.g., via session factories)
+    - Services are used in contexts with different session lifecycle requirements
+
+    Args:
+        db_session: Optional database session to use immediately. If not provided,
+            the session must be set later via `set_db_session()` or a factory must be provided.
+        db_session_factory: Optional callable that returns a database session when invoked.
+            Used as a fallback if no session is directly provided.
+
+    Raises:
+        RuntimeError: When accessing `db_session` property without a session or factory configured.
+
+    Example:
+        >>> # With immediate session
+        >>> service = MyService(db_session=session)
+        >>>
+        >>> # With deferred session
+        >>> service = MyService()
+        >>> service.set_db_session(session)
+        >>>
+        >>> # With session factory
+        >>> service = MyService(db_session_factory=lambda: get_session())
+    """
+
+    def __init__(
+        self,
+        db_session: Session | None = None,
+        db_session_factory: Callable[[], Session] | None = None,
+    ):
+        self._db_session: Session | None = db_session
+        self._db_session_factory = db_session_factory
+        self._session_managed_services: list[BaseSessionManagedService] = []
+
+    def set_db_session(self, db_session: Session) -> None:
+        """Set the database session for the service."""
+        self._db_session = db_session
+        for service in self._session_managed_services:
+            service.set_db_session(db_session)
+
+    def register_managed_services(self, *services: "BaseSessionManagedService") -> None:
+        """Register a child service that also requires session management."""
+        self._session_managed_services.extend(services)
+
+    @property
+    def db_session(self) -> Session:
+        if self._db_session is not None:
+            return self._db_session
+        if self._db_session_factory is not None:
+            return self._db_session_factory()
+        raise RuntimeError("No DB session available. Provide session or session factory.")

application/backend/app/services/data_collect/data_collector.py

@@ -180,10 +180,11 @@ def collect(
             return
         frame_data = cv2.cvtColor(frame_data, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
         with get_db_session() as session:
-            labels = LabelService(db_session=session).list_all(project_id=project.id)
+            label_service = LabelService(db_session=session)
+            labels = label_service.list_all(project_id=project.id)
             annotations = convert_prediction(labels=labels, frame_data=frame_data, prediction=inference_data.prediction)
 
-            dataset_service = DatasetService(data_dir=self.data_dir, db_session=session)
+            dataset_service = DatasetService(data_dir=self.data_dir, label_service=label_service, db_session=session)
             dataset_service.create_dataset_item(
                 project=project,
                 name=f"{timestamp:.4f}".replace(".", "_"),