Skip to content

Commit 03605b0

Browse files
committed
update to more closely relate to project.export_labels()
1 parent 98075b1 commit 03605b0

File tree

1 file changed

+28
-33
lines changed

1 file changed

+28
-33
lines changed

labelbox/schema/model_run.py

Lines changed: 28 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import TYPE_CHECKING, Dict, Iterable, Union
1+
from typing import TYPE_CHECKING, Dict, Iterable, Union, List, Optional, Any
22
from pathlib import Path
33
import os
44
import time
@@ -182,7 +182,11 @@ def delete_model_run_data_rows(self, data_row_ids):
182182
})
183183

184184
@experimental
185-
def export_annotations(self, download: bool = False) -> LabelGenerator:
185+
def export_labels(
186+
self,
187+
download: bool = False,
188+
timeout_seconds: int = 600
189+
) -> Optional[Union[str, List[Dict[Any, Any]]]]:
186190
"""
187191
Experimental. To use, make sure client has enable_experimental=True.
188192
@@ -191,49 +195,40 @@ def export_annotations(self, download: bool = False) -> LabelGenerator:
191195
Args:
192196
download (bool): Returns the url if False
193197
Returns:
194-
URL of the data file that is generated from this ModelRun.
198+
URL of the data file with this ModelRun's labels.
195199
If download=True, this instead returns the contents as NDJSON format.
200+
If the server didn't generate during the `timeout_seconds` period,
201+
None is returned.
196202
"""
197-
203+
sleep_time = 2
198204
query_str = """
199205
mutation exportModelRunAnnotationsPyApi($modelRunId: ID!) {
200206
exportModelRunAnnotations(data: {modelRunId: $modelRunId}) {
201207
downloadUrl createdAt status
202208
}
203209
}
204210
"""
205-
url = self.client.execute(
206-
query_str, {'modelRunId': self.model_id},
207-
experimental=True)['exportModelRunAnnotations']['downloadUrl']
208-
209-
counter = 1
210-
while url is None:
211-
logger.info(f"Fetching model run annotations...attempt: {counter}")
212-
counter += 1
213-
if counter == 10:
214-
raise Exception(
215-
f"Unsuccessfully got downloadUrl after {counter} attempts.")
216-
time.sleep(10)
211+
212+
while True:
217213
url = self.client.execute(
218-
query_str, {'modelRunId': self.model_id},
214+
query_str, {'modelRunId': self.uid},
219215
experimental=True)['exportModelRunAnnotations']['downloadUrl']
220216

221-
if not download:
222-
return url
223-
else:
224-
response = requests.get(url)
225-
response.raise_for_status()
226-
contents = ndjson.loads(response.content)
227-
return contents
228-
"""
229-
# Need to discuss how to properly get the asset data type from export
230-
# Until then, we can't do much with contents as the LBV1Converter may
231-
# not infer correct media type and error out
232-
# for now, return the contents as either url or ndjson content
233-
# """
234-
235-
# labels = LBV1Converter.deserialize(contents)
236-
# return labels
217+
if url:
218+
if not download:
219+
return url
220+
else:
221+
response = requests.get(url)
222+
response.raise_for_status()
223+
return ndjson.loads(response.content)
224+
225+
timeout_seconds -= sleep_time
226+
if timeout_seconds <= 0:
227+
return None
228+
229+
logger.debug("ModelRun '%s' label export, waiting for server...",
230+
self.uid)
231+
time.sleep(sleep_time)
237232

238233

239234
class ModelRunDataRow(DbObject):

0 commit comments

Comments
 (0)