@@ -50,7 +50,7 @@ def test_repr(test_object: any):
5050def model_run (CLIENT ):
5151 ds = CLIENT .create_dataset (TEST_DATASET_NAME )
5252 ds_items = []
53- for url in TEST_IMG_URLS :
53+ for url in TEST_IMG_URLS [: 2 ] :
5454 ds_items .append (
5555 DatasetItem (
5656 image_location = url ,
@@ -246,7 +246,7 @@ def test_polygon_pred_upload_ignore(model_run):
246246 )
247247
248248
249- def test_mixed_pred_upload (model_run ):
249+ def test_mixed_pred_upload (model_run : ModelRun ):
250250 prediction_semseg = SegmentationPrediction .from_json (
251251 TEST_SEGMENTATION_PREDICTIONS [0 ]
252252 )
@@ -262,15 +262,15 @@ def test_mixed_pred_upload(model_run):
262262 assert response ["predictions_processed" ] == 3
263263 assert response ["predictions_ignored" ] == 0
264264
265- response_refloc = model_run .refloc ( prediction_polygon . reference_id )
265+ all_predictions = model_run .ungrouped_export ( )
266266 assert_box_prediction_matches_dict (
267- response_refloc ["box" ][0 ], TEST_BOX_PREDICTIONS [0 ]
267+ all_predictions ["box" ][0 ], TEST_BOX_PREDICTIONS [0 ]
268268 )
269269 assert_polygon_prediction_matches_dict (
270- response_refloc ["polygon" ][0 ], TEST_POLYGON_PREDICTIONS [0 ]
270+ all_predictions ["polygon" ][0 ], TEST_POLYGON_PREDICTIONS [0 ]
271271 )
272272 assert_segmentation_annotation_matches_dict (
273- response_refloc ["segmentation" ][0 ], TEST_SEGMENTATION_PREDICTIONS [0 ]
273+ all_predictions ["segmentation" ][0 ], TEST_SEGMENTATION_PREDICTIONS [0 ]
274274 )
275275
276276
0 commit comments