diff --git a/libs/labelbox/src/labelbox/schema/ontology.py b/libs/labelbox/src/labelbox/schema/ontology.py index 3238dc4a5..1cc827c5f 100644 --- a/libs/labelbox/src/labelbox/schema/ontology.py +++ b/libs/labelbox/src/labelbox/schema/ontology.py @@ -30,6 +30,7 @@ FeatureSchemaAttributes, ) import warnings +from labelbox.schema.project import MediaType class DeleteFeatureFromOntologyResult: @@ -195,6 +196,7 @@ class Ontology(DbObject): normalized = Field.Json("normalized") object_schema_count = Field.Int("object_schema_count") classification_schema_count = Field.Int("classification_schema_count") + media_type = Field.Enum(MediaType, "media_type", "mediaType") projects = Relationship.ToMany("Project", True) created_by = Relationship.ToOne("User", False, "created_by") diff --git a/libs/labelbox/src/labelbox/schema/project.py b/libs/labelbox/src/labelbox/schema/project.py index 9494d33cf..c30db69b7 100644 --- a/libs/labelbox/src/labelbox/schema/project.py +++ b/libs/labelbox/src/labelbox/schema/project.py @@ -666,14 +666,30 @@ def review_metrics(self, net_score) -> int: def connect_ontology(self, ontology) -> None: """ Connects the ontology to the project. If an editor is not setup, it will be connected as well. + This method can be used to change the project's ontology. Note: For live chat model evaluation projects, the editor setup is skipped because it is automatically setup when the project is created. Args: ontology (Ontology): The ontology to attach to the project + + Raises: + ValueError: If ontology and project have different media types and ontology has a media type set """ - if not self.is_empty_ontology(): - raise ValueError("Ontology already connected to project.") + # Check media type compatibility + if ( + self.media_type != ontology.media_type + and not ontology.media_type == MediaType.Unknown + ): + raise ValueError( + "Ontology and project must share the same type, unless the ontology has no type." + ) + + # Check if project has labels and warn user + if self.get_label_count() > 0: + warnings.warn( + "Project has labels. The new ontology must contain all annotation types." + ) if ( self.labeling_frontend() is None diff --git a/libs/labelbox/tests/integration/test_project_setup.py b/libs/labelbox/tests/integration/test_project_setup.py index 31cf5e066..1a5091f3a 100644 --- a/libs/labelbox/tests/integration/test_project_setup.py +++ b/libs/labelbox/tests/integration/test_project_setup.py @@ -2,6 +2,7 @@ from datetime import datetime, timedelta, timezone import pytest +from labelbox.schema.media_type import MediaType def simple_ontology(): @@ -38,11 +39,69 @@ def test_project_editor_setup(client, project, rand_gen): ] == [ontology_name] -def test_project_connect_ontology_cant_call_multiple_times( - client, project, rand_gen -): - ontology_name = f"test_project_editor_setup_ontology_name-{rand_gen(str)}" - ontology = client.create_ontology(ontology_name, simple_ontology()) - project.connect_ontology(ontology) - with pytest.raises(ValueError): - project.connect_ontology(ontology) +def test_project_connect_ontology_multiple_times(client, project, rand_gen): + """Test that we can connect multiple ontologies in sequence.""" + # Connect first ontology + ontology_name_1 = ( + f"test_project_connect_ontology_multiple_times_1-{rand_gen(str)}" + ) + ontology_1 = client.create_ontology(ontology_name_1, simple_ontology()) + project.connect_ontology(ontology_1) + assert project.ontology().name == ontology_name_1 + + # Connect second ontology + ontology_name_2 = ( + f"test_project_connect_ontology_multiple_times_2-{rand_gen(str)}" + ) + ontology_2 = client.create_ontology(ontology_name_2, simple_ontology()) + project.connect_ontology(ontology_2) + assert project.ontology().name == ontology_name_2 + + +def test_project_connect_ontology_with_different_media_types(client, rand_gen): + """Test connecting ontologies with different media types to a project""" + # Create a new project with Image media type + project_name = f"test_project_media_type_{rand_gen(str)}" + project = client.create_project( + name=project_name, media_type=MediaType.Image + ) + + try: + # Create ontologies with different media types + ontology_1 = client.create_ontology( + f"test_ontology_1_{rand_gen(str)}", + simple_ontology(), + media_type=MediaType.Image, # Same media type as project + ) + + ontology_2 = client.create_ontology( + f"test_ontology_2_{rand_gen(str)}", + simple_ontology(), + media_type=MediaType.Video, # Different media type + ) + + # Test connecting ontology with same media type + project.connect_ontology(ontology_1) + assert project.ontology().uid == ontology_1.uid + + # Test connecting ontology with different media type + with pytest.raises(ValueError) as exc_info: + project.connect_ontology(ontology_2) + assert "Ontology and project must share the same type" in str( + exc_info.value + ) + finally: + # Clean up + project.delete() + + +def test_project_connect_ontology_with_unknown_type(client, project, rand_gen): + """Test connecting ontology with unknown media type to a project""" + # Create ontology with unknown media type + unknown_type_ontology = client.create_ontology( + f"test_unknown_type_{rand_gen(str)}", simple_ontology() + ) + + # Test connecting ontology with unknown type + project.connect_ontology(unknown_type_ontology) + assert project.ontology().uid == unknown_type_ontology.uid