Skip to content

Commit c5865c0

Browse files
committed
WIP - working on creating an export for annotations on a model run
1 parent ac1a037 commit c5865c0

File tree

1 file changed

+62
-1
lines changed

1 file changed

+62
-1
lines changed

labelbox/schema/model_run.py

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,21 @@
22
from pathlib import Path
33
import os
44
import time
5+
import logging
6+
import requests
7+
import ndjson
8+
from labelbox.data.annotation_types.collection import LabelGenerator
59

610
from labelbox.pagination import PaginatedCollection
711
from labelbox.orm.query import results_query_part
812
from labelbox.orm.model import Field, Relationship, Entity
9-
from labelbox.orm.db_object import DbObject
13+
from labelbox.orm.db_object import DbObject, experimental
1014

1115
if TYPE_CHECKING:
1216
from labelbox import MEAPredictionImport
1317

18+
logger = logging.getLogger(__name__)
19+
1420

1521
class ModelRun(DbObject):
1622
name = Field.String("name")
@@ -175,6 +181,61 @@ def delete_model_run_data_rows(self, data_row_ids):
175181
data_row_ids_param: data_row_ids
176182
})
177183

184+
@experimental
185+
def export_model_run_annotations(self,
186+
download: bool = False) -> LabelGenerator:
187+
"""
188+
Experimental. To use, make sure client has enable_experimental=True.
189+
190+
Fetches Labels from the ModelRun
191+
192+
Args:
193+
download (bool): Returns the url if False
194+
Returns:
195+
URL of the data file that is generated from this ModelRun.
196+
If download=True, this instead returns the contents as NDJSON format.
197+
"""
198+
199+
query_str = """
200+
mutation exportModelRunAnnotationsPyApi($modelRunId: ID!) {
201+
exportModelRunAnnotations(data: {modelRunId: $modelRunId}) {
202+
downloadUrl createdAt status
203+
}
204+
}
205+
"""
206+
url = self.client.execute(
207+
query_str, {'modelRunId': self.model_id},
208+
experimental=True)['exportModelRunAnnotations']['downloadUrl']
209+
210+
counter = 1
211+
while url is None:
212+
logger.info(f"Fetching model run annotations...attempt: {counter}")
213+
counter += 1
214+
if counter == 10:
215+
raise Exception(
216+
f"Unsuccessfully got downloadUrl after {counter} attempts.")
217+
time.sleep(10)
218+
url = self.client.execute(
219+
query_str, {'modelRunId': self.model_id},
220+
experimental=True)['exportModelRunAnnotations']['downloadUrl']
221+
222+
if not download:
223+
return url
224+
else:
225+
response = requests.get(url)
226+
response.raise_for_status()
227+
contents = ndjson.loads(response.content)
228+
return contents
229+
"""
230+
# Need to discuss how to properly get the asset data type from export
231+
# Until then, we can't do much with contents as the LBV1Converter may
232+
# not infer correct media type and error out
233+
# for now, return the contents as either url or ndjson content
234+
# """
235+
236+
# labels = LBV1Converter.deserialize(contents)
237+
# return labels
238+
178239

179240
class ModelRunDataRow(DbObject):
180241
label_id = Field.String("label_id")

0 commit comments

Comments
 (0)