Store training dataset (#5008)

Author: itallix

application/backend/app/repositories/__init__.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .dataset_item_repo import DatasetItemRepository
+from .dataset_revision_repo import DatasetRevisionRepository
 from .label_repo import LabelRepository
 from .model_revision_repo import ModelRevisionRepository
 from .pipeline_repo import PipelineRepository
@@ -11,6 +12,7 @@
 
 __all__ = [
     "DatasetItemRepository",
+    "DatasetRevisionRepository",
     "LabelRepository",
     "ModelRevisionRepository",
     "PipelineRepository",

application/backend/app/repositories/dataset_revision_repo.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from sqlalchemy.orm import Session
+
+from app.db.schema import DatasetRevisionDB
+from app.repositories.base import BaseRepository
+
+
+class DatasetRevisionRepository(BaseRepository[DatasetRevisionDB]):
+    """Repository for dataset revision-related database operations."""
+
+    def __init__(self, db: Session):
+        super().__init__(db, DatasetRevisionDB)

application/backend/app/services/dataset_service.py

@@ -13,11 +13,12 @@
 
 import datumaro.experimental as dm
 import numpy as np
+from datumaro.experimental.export_import import export_dataset
 from loguru import logger
 from PIL import Image, UnidentifiedImageError
 from sqlalchemy.orm import Session
 
-from app.db.schema import DatasetItemDB
+from app.db.schema import DatasetItemDB, DatasetRevisionDB
 from app.models import (
     DatasetItem,
     DatasetItemAnnotation,
@@ -29,7 +30,7 @@
     Rectangle,
     TaskType,
 )
-from app.repositories import DatasetItemRepository
+from app.repositories import DatasetItemRepository, DatasetRevisionRepository
 from app.schemas.project import ProjectBase, ProjectView, TaskBase
 from app.services.datumaro_converter import convert_dataset
 from app.utils.images import crop_to_thumbnail
@@ -373,3 +374,32 @@ def _get_image_path(item: DatasetItem) -> str:
             get_dataset_items=_get_dataset_items,
             get_image_path=_get_image_path,
         )
+
+    def save_revision(self, project_id: UUID, dataset: dm.Dataset) -> None:
+        """
+        Saves the dataset as a new revision.
+
+        Creates a new dataset revision entry in the database and exports the dataset
+        to a zip file in the project's revisions directory.
+
+        Args:
+            project_id: The UUID of the project to save the revision for.
+            dataset: The Datumaro dataset to export.
+
+        Returns:
+            None
+        """
+        revision_repo = DatasetRevisionRepository(db=self.db_session)
+        revision_db = revision_repo.save(
+            DatasetRevisionDB(
+                project_id=str(project_id),
+            )
+        )
+        revision_path = self.projects_dir / str(project_id) / "dataset_revisions" / revision_db.id
+        logger.info("Saving dataset revision '{}' to '{}'", revision_db.id, revision_path)
+        export_dataset(
+            dataset=dataset,
+            output_path=revision_path,
+            export_images=True,
+            as_zip=True,
+        )

application/backend/app/services/training/otx_trainer.py

@@ -127,6 +127,7 @@ def create_training_dataset(self) -> None:
             self._training_dataset = dm_dataset.filter_by_subset(Subset.TRAINING)
             self._validation_dataset = dm_dataset.filter_by_subset(Subset.VALIDATION)
             self._testing_dataset = dm_dataset.filter_by_subset(Subset.TESTING)
+            self._dataset_service.save_revision(project_id, dm_dataset)
 
     @step("Train Model with OTX")
     def train_model(self) -> None:

application/backend/tests/integration/services/test_dataset_service.py

@@ -12,7 +12,7 @@
 from sqlalchemy import func, select
 from sqlalchemy.orm import Session
 
-from app.db.schema import DatasetItemDB, DatasetItemLabelDB, PipelineDB
+from app.db.schema import DatasetItemDB, DatasetItemLabelDB, DatasetRevisionDB, PipelineDB
 from app.models import DatasetItemAnnotation, DatasetItemAnnotationStatus, DatasetItemSubset, LabelReference, Rectangle
 from app.schemas import PipelineView, ProjectView
 from app.services import LabelService, PipelineService, ProjectService
@@ -1373,3 +1373,25 @@ def test_subset_filter_verifies_data_correctness(
         assert len(testing_items) == 1
         for item in testing_items:
             assert item.subset == DatasetItemSubset.TESTING
+
+    def test_save_revision(
+        self,
+        fxt_projects_dir: Path,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_subset_items: tuple[ProjectView, list[DatasetItemDB]],
+        db_session: Session,
+    ) -> None:
+        """Test saving a dataset revision."""
+        project, db_dataset_items = fxt_project_with_subset_items
+        dataset = fxt_dataset_service.get_dm_dataset(project.id, project.task, DatasetItemAnnotationStatus.REVIEWED)
+
+        fxt_dataset_service.save_revision(
+            project_id=project.id,
+            dataset=dataset,
+        )
+
+        # Verify that a revision entry was created
+        db_revisions = db_session.query(DatasetRevisionDB).all()
+        assert len(db_revisions) == 1
+        revision_id = db_revisions[0].id
+        assert (fxt_projects_dir / str(project.id) / "dataset_revisions" / revision_id / "dataset.zip").exists()

application/backend/tests/unit/services/training/test_otx_trainer.py

@@ -319,6 +319,7 @@ def test_create_training_dataset_success(
         assert otx_trainer._training_dataset == mock_training_dataset
         assert otx_trainer._validation_dataset == mock_validation_dataset
         assert otx_trainer._testing_dataset == mock_testing_dataset
+        fxt_dataset_service.save_revision.assert_called_once_with(project_id, mock_dm_dataset)
 
     def test_create_training_dataset_without_project_id_raises_error(
         self,