Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions labelbox/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@
from labelbox.schema.role import Role, ProjectRole
from labelbox.schema.invite import Invite, InviteLimit
from labelbox.schema.model_run import ModelRun
from labelbox.schema.model import Model
16 changes: 16 additions & 0 deletions labelbox/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -533,3 +533,19 @@ def get_models(self, where=None):
An iterable of Models (typically a PaginatedCollection).
"""
return self._get_all(Model, where, filter_deleted=False)

def create_model(self, name, ontology_id):
""" Creates a Model object on the server.

>>> model = client.create_model(<model_name>, <ontology_id>)

Args:
name (string): Name of the model
ontology_id (string): ID of the related ontology
Returns:
A new Model object.
Raises:
InvalidAttributeError: If the Model type does not contain
any of the attribute names given in kwargs.
"""
return self._create(Model, {"name": name, "ontology_id": ontology_id})
11 changes: 6 additions & 5 deletions labelbox/schema/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,17 @@
import labelbox.schema.benchmark
import labelbox.schema.data_row
import labelbox.schema.dataset
import labelbox.schema.invite
import labelbox.schema.label
import labelbox.schema.labeling_frontend
import labelbox.schema.model
import labelbox.schema.model_run
import labelbox.schema.ontology
import labelbox.schema.organization
import labelbox.schema.prediction
import labelbox.schema.project
import labelbox.schema.review
import labelbox.schema.role
import labelbox.schema.task
import labelbox.schema.user
import labelbox.schema.webhook
import labelbox.schema.prediction
import labelbox.schema.ontology
import labelbox.schema.invite
import labelbox.schema.role
import labelbox.schema.model_run
26 changes: 25 additions & 1 deletion labelbox/schema/model.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from labelbox.orm import query
from labelbox.orm.db_object import DbObject
from labelbox.orm.model import Field, Relationship
from labelbox.orm.model import Entity, Field, Relationship


class Model(DbObject):
Expand All @@ -16,3 +17,26 @@ class Model(DbObject):

model_runs = Relationship.ToMany("ModelRun", False)
ontology = Relationship.ToOne("Ontology", False)

def create_model_run(self, name):
""" Creates a model run belonging to this model.

Args:
name (string): The name for the model run.
Returns:
ModelRun, the created model run.
"""
name_param = "name"
model_id_param = "modelId"
ModelRun = Entity.ModelRun
query_str = """mutation CreateModelRunPyApi($%s: String!, $%s: ID!) {
createModelRun(data: {name: $%s, modelId: $%s}) {%s}}""" % (
name_param, model_id_param,
name_param, model_id_param,
query.results_query_part(ModelRun)
)
res = self.client.execute(query_str, {
name_param: name,
model_id_param: self.uid
})
return ModelRun(self.client, res["createModelRun"])
23 changes: 23 additions & 0 deletions tests/integration/test_model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
from labelbox import Model

def test_model(client, configured_project, rand_gen):
before = list(client.get_models())
for m in before:
assert isinstance(m, Model)

ontology = configured_project.ontology

data = {"name": rand_gen(str), "ontology_id": ontology.uid}
model = client.create_model(data["name"], data["ontology_id"])
assert model.name == data["name"]
assert model.ontology.id == data["ontology_id"]

after = list(client.get_models())
assert len(after) == len(before) + 1
assert model in after

model = client.get_model(model.uid)
assert model.name == data["name"]
assert model.ontology.id == data["ontology_id"]