def __enter__(self) -> resources_pb2.Input:
my_concept_id = "my-concept-id-" + uuid.uuid4().hex
my_concept_name = "my concept name " + uuid.uuid4().hex
image_metadata = struct_pb2.Struct()
image_metadata.update({
"some-key": "some-value",
"another-key": {
"inner-key": "inner-value"
}
})
post_response = self._stub.PostInputs(
service_pb2.PostInputsRequest(inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL,
allow_duplicate_url=True),
concepts=[
resources_pb2.Concept(
id=my_concept_id, name=my_concept_name, value=1)
],
metadata=image_metadata,
geo=resources_pb2.Geo(geo_point=resources_pb2.GeoPoint(
longitude=44, latitude=55)),
), )
]),
metadata=metadata(),
)
raise_on_failure(post_response)
self._input = post_response.inputs[0]
wait_for_inputs_upload(self._stub, metadata(), [self._input.id])
return self._input
def test_post_delete_batch_images(channel):
stub = service_pb2_grpc.V2Stub(channel)
post_response = stub.PostInputs(
service_pb2.PostInputsRequest(
inputs=[
resources_pb2.Input(
data=resources_pb2.Data(
image=resources_pb2.Image(url=TRUCK_IMAGE_URL, allow_duplicate_url=True)
)
),
resources_pb2.Input(
data=resources_pb2.Data(
image=resources_pb2.Image(url=TRUCK_IMAGE_URL, allow_duplicate_url=True)
)
),
]
),
metadata=metadata(),
)
raise_on_failure(post_response)
input_id1 = post_response.inputs[0].id
input_id2 = post_response.inputs[1].id
wait_for_inputs_upload(stub, metadata(), [input_id1, input_id2])
delete_response = stub.DeleteInputs(
service_pb2.DeleteInputsRequest(ids=[input_id1, input_id2]), metadata=metadata()
)
raise_on_failure(delete_response)
def test_save_and_execute_search_by_id(channel):
stub = service_pb2_grpc.V2Stub(channel)
search_id = "my-search-id-" + uuid.uuid4().hex[:15]
with SetupImage(stub) as input_:
my_concept_id = input_.data.concepts[0].id
# This saves the search under an ID, but does not execute it / return any results.
save_search_response = stub.PostSearches(
service_pb2.PostSearchesRequest(searches=[
resources_pb2.Search(
id=search_id,
save=True,
query=resources_pb2.Query(ands=[
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(id=my_concept_id,
value=1)
])))
]),
)
]),
metadata=metadata(),
)
raise_on_failure(save_search_response)
# Executing the search returns results.
post_search_by_id_response = stub.PostSearchesByID(
service_pb2.PostSearchesByIDRequest(id=search_id),
metadata=metadata(),
)
raise_on_failure(post_search_by_id_response)
assert len(post_search_by_id_response.hits) == 1
assert post_search_by_id_response.hits[0].input.id == input_.id
def test_post_model_with_hyper_params(channel):
stub = service_pb2_grpc.V2Stub(channel)
model_id = uuid.uuid4().hex[:30]
hyper_params = struct_pb2.Struct()
hyper_params.update({
"MAX_NITEMS": 1000000,
"MIN_NITEMS": 1000,
"N_EPOCHS": 5,
"custom_training_cfg": "custom_training_1layer",
"custom_training_cfg_args": {},
})
post_response = stub.PostModels(
service_pb2.PostModelsRequest(models=[
resources_pb2.Model(
id=model_id,
output_info=resources_pb2.OutputInfo(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(id="some-initial-concept")
], ),
output_config=resources_pb2.OutputConfig(
hyper_params=hyper_params),
),
)
]),
metadata=metadata(),
)
raise_on_failure(post_response)
assert (post_response.model.output_info.output_config.
hyper_params["custom_training_cfg"] == "custom_training_1layer")
delete_response = stub.DeleteModel(
service_pb2.DeleteModelRequest(model_id=model_id), metadata=metadata())
raise_on_failure(delete_response)
def test_image_with_bytes(channel):
stub = service_pb2_grpc.V2Stub(channel)
with open(RED_TRUCK_IMAGE_FILE_PATH, "rb") as f:
file_bytes = f.read()
post_response = stub.PostInputs(
service_pb2.PostInputsRequest(
inputs=[
resources_pb2.Input(
data=resources_pb2.Data(image=resources_pb2.Image(base64=file_bytes))
)
]
),
metadata=metadata(),
)
raise_on_failure(post_response)
input_id = post_response.inputs[0].id
wait_for_inputs_upload(stub, metadata(), [input_id])
delete_response = stub.DeleteInputs(
service_pb2.DeleteInputsRequest(ids=[input_id]), metadata=metadata()
)
raise_on_failure(delete_response)
def test_concept_post_get_patch(channel):
stub = service_pb2_grpc.V2Stub(channel)
random_string = uuid.uuid4().hex
random_concept_id = "some-concept-id-狗-" + random_string
random_concept_name = "some-concept-name-的な-" + random_string
post_concepts_response = stub.PostConcepts(
service_pb2.PostConceptsRequest(concepts=[
resources_pb2.Concept(id=random_concept_id,
name=random_concept_name)
]),
metadata=metadata(),
)
raise_on_failure(post_concepts_response)
get_concepts_response = stub.GetConcept(
service_pb2.GetConceptRequest(concept_id=random_concept_id),
metadata=metadata())
raise_on_failure(get_concepts_response)
assert get_concepts_response.concept.id == random_concept_id
assert get_concepts_response.concept.name == random_concept_name
duplicated_post_concepts_response = stub.PostConcepts(
service_pb2.PostConceptsRequest(
concepts=[resources_pb2.Concept(id=random_concept_id, )]),
metadata=metadata(),
)
assert (duplicated_post_concepts_response.status.code ==
status_code_pb2.StatusCode.CONCEPTS_INVALID_REQUEST)
assert duplicated_post_concepts_response.status.description == "Invalid request"
assert "duplicate" in duplicated_post_concepts_response.status.details.lower(
)
post_concepts_searches_response = stub.PostConceptsSearches(
service_pb2.PostConceptsSearchesRequest(
concept_query=resources_pb2.ConceptQuery(
name=random_concept_name)),
metadata=metadata(),
)
raise_on_failure(post_concepts_searches_response)
assert random_concept_name in post_concepts_searches_response.concepts[
0].name
patch_concepts_response = stub.PatchConcepts(
service_pb2.PatchConceptsRequest(
action="overwrite",
concepts=[
resources_pb2.Concept(id=random_concept_id,
name="some new concept name")
],
),
metadata=metadata(),
)
raise_on_failure(patch_concepts_response)
def test_list_models_with_pagination(channel):
stub = service_pb2_grpc.V2Stub(channel)
response = stub.ListModels(service_pb2.ListModelsRequest(per_page=2), metadata=metadata())
raise_on_failure(response)
assert len(response.models) == 2
# We shouldn't have 1000*500 number of models, so the result should be empty.
response = stub.ListModels(
service_pb2.ListModelsRequest(page=1000, per_page=500), metadata=metadata()
)
raise_on_failure(response)
assert len(response.models) == 0
def __enter__(self) -> resources_pb2.Input:
post_response = self._stub.PostInputs(
service_pb2.PostInputsRequest(inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=TRAVEL_IMAGE_URL,
allow_duplicate_url=True), ), )
]),
metadata=metadata(),
)
raise_on_failure(post_response)
self._input = post_response.inputs[0]
wait_for_inputs_upload(self._stub, metadata(), [self._input.id])
return self._input
def test_list_all_models(channel):
stub = service_pb2_grpc.V2Stub(channel)
list_response = stub.ListModels(service_pb2.ListModelsRequest(),
metadata=metadata())
raise_on_failure(list_response)
assert len(list_response.models) > 0
def test_predict_video_url_with_custom_sample_ms(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
video=resources_pb2.Video(url=BEER_VIDEO_URL)))
],
model=resources_pb2.Model(output_info=resources_pb2.OutputInfo(
output_config=resources_pb2.OutputConfig(sample_ms=2000))),
)
response = post_model_outputs_and_maybe_allow_retries(stub,
request,
metadata=metadata())
raise_on_failure(response)
# The expected time per frame is the middle between the start and the end of the frame
# (in milliseconds).
expected_time = 1000
assert len(response.outputs[0].data.frames) > 0
for frame in response.outputs[0].data.frames:
assert frame.frame_info.time == expected_time
expected_time += 2000
def test_search_by_geo_box_and_annotated_name_and_predicted_name(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
my_concept_name = input_.data.concepts[0].name
response = stub.PostSearches(
service_pb2.PostSearchesRequest(
query=resources_pb2.Query(ands=[
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(geo=resources_pb2.Geo(geo_box=[
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(longitude=43,
latitude=54)),
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(longitude=45,
latitude=56)),
])))),
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(name=my_concept_name,
value=1)
]))),
resources_pb2.And(output=resources_pb2.Output(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(name="dog", value=1)
]))),
]),
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) > 0
assert input_.id in [hit.input.id for hit in response.hits]
def test_search_by_image_url_and_geo_box(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
response = stub.PostInputsSearches(
PostInputsSearchesRequest(
searches=[
Search(query=Query(ranks=[
Rank(annotation=Annotation(data=Data(image=Image(
url=DOG_IMAGE_URL)))),
Rank(annotation=Annotation(data=Data(
geo=resources_pb2.Geo(geo_box=[
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(
longitude=43, latitude=54)),
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(
longitude=45, latitude=56)),
])))),
]))
],
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) >= 1
assert input_.id in [hit.input.id for hit in response.hits]
def test_predict_image_url_with_selected_concepts(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL, ), ), )
],
model=resources_pb2.Model(output_info=resources_pb2.OutputInfo(
output_config=resources_pb2.OutputConfig(select_concepts=[
resources_pb2.Concept(name="dog"),
resources_pb2.Concept(name="cat"),
]))),
)
response = post_model_outputs_and_maybe_allow_retries(stub,
request,
metadata=metadata())
raise_on_failure(response)
concepts = response.outputs[0].data.concepts
assert len(concepts) == 2
dog_concept = [c for c in concepts if c.name == "dog"][0]
cat_concept = [c for c in concepts if c.name == "cat"][0]
assert dog_concept.value > cat_concept.value
def test_search_by_image_url_and_geo_box(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
response = stub.PostSearches(
service_pb2.PostSearchesRequest(
query=resources_pb2.Query(ands=[
resources_pb2.And(output=resources_pb2.Output(
input=resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL))))),
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(geo=resources_pb2.Geo(geo_box=[
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(longitude=43,
latitude=54)),
resources_pb2.GeoBoxedPoint(
geo_point=resources_pb2.GeoPoint(longitude=45,
latitude=56)),
])))),
]),
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) > 0
assert input_.id in [hit.input.id for hit in response.hits]
def test_post_list_patch_get_delete_image(channel):
stub = service_pb2_grpc.V2Stub(channel)
post_response = stub.PostInputs(
service_pb2.PostInputsRequest(
inputs=[
resources_pb2.Input(
data=resources_pb2.Data(
image=resources_pb2.Image(url=TRUCK_IMAGE_URL, allow_duplicate_url=True),
concepts=[resources_pb2.Concept(id="some-concept")],
)
)
]
),
metadata=metadata(),
)
raise_on_failure(post_response)
input_id = post_response.inputs[0].id
try:
wait_for_inputs_upload(stub, metadata(), [input_id])
list_response = stub.ListInputs(
service_pb2.ListInputsRequest(per_page=1), metadata=metadata()
)
raise_on_failure(list_response)
assert len(list_response.inputs) == 1
# Most likely we don"t have that many inputs, so this should return 0.
list_response2 = stub.ListInputs(
service_pb2.ListInputsRequest(per_page=500, page=1000), metadata=metadata()
)
raise_on_failure(list_response2)
assert len(list_response2.inputs) == 0
patch_response = stub.PatchInputs(
service_pb2.PatchInputsRequest(
action="overwrite",
inputs=[
resources_pb2.Input(
id=input_id,
data=resources_pb2.Data(
concepts=[resources_pb2.Concept(id="some-new-concept")]
),
)
],
),
metadata=metadata(),
)
raise_on_failure(patch_response)
get_response = stub.GetInput(
service_pb2.GetInputRequest(input_id=input_id), metadata=metadata()
)
raise_on_failure(get_response)
assert get_response.input.data.concepts[0].name == "some-new-concept"
finally:
delete_request = service_pb2.DeleteInputRequest(input_id=input_id)
delete_response = stub.DeleteInput(delete_request, metadata=metadata())
raise_on_failure(delete_response)
def test_search_public_concepts_in_english(channel):
stub = service_pb2_grpc.V2Stub(channel)
post_concepts_searches_response = stub.PostConceptsSearches(
service_pb2.PostConceptsSearchesRequest(
concept_query=resources_pb2.ConceptQuery(name="dog*")
),
metadata=metadata(),
)
raise_on_failure(post_concepts_searches_response)
assert len(post_concepts_searches_response.concepts) > 0
def test_search_for_model(channel):
stub = service_pb2_grpc.V2Stub(channel)
response = stub.PostModelsSearches(
service_pb2.PostModelsSearchesRequest(
model_query=resources_pb2.ModelQuery(name="*general*")),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.models) > 0
for m in response.models:
assert "general" in m.name.lower()
def test_predict_image_url(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL)))
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
raise_on_failure(response)
assert len(response.outputs[0].data.concepts) > 0
def test_failed_predict(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=NON_EXISTING_IMAGE_URL)))
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
assert response.status.code == status_code_pb2.FAILURE
assert response.status.description == "Failure"
assert response.outputs[
0].status.code == status_code_pb2.INPUT_DOWNLOAD_FAILED
def test_predict_video_url(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
video=resources_pb2.Video(url=CONAN_GIF_VIDEO_URL)))
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
raise_on_failure(response)
assert len(response.outputs[0].data.frames) > 0
for frame in response.outputs[0].data.frames:
assert len(frame.data.concepts) > 0
def test_predict_image_url_with_max_concepts(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL, ), ), )
],
model=resources_pb2.Model(output_info=resources_pb2.OutputInfo(
output_config=resources_pb2.OutputConfig(max_concepts=3))),
)
response = stub.PostModelOutputs(request, metadata=metadata())
raise_on_failure(response)
assert len(response.outputs[0].data.concepts) == 3
def test_search_by_annotated_concept_name(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
my_concept_name = input_.data.concepts[0].name
response = stub.PostSearches(
service_pb2.PostSearchesRequest(query=resources_pb2.Query(ands=[
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(name=my_concept_name, value=1)
])))
])),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) == 1
assert response.hits[0].input.id == input_.id
def test_image_predict_on_public_models(channel):
stub = service_pb2_grpc.V2Stub(channel)
for title, model_id in MODEL_TITLE_AND_ID_PAIRS:
request = service_pb2.PostModelOutputsRequest(
model_id=model_id,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL)))
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
raise_on_failure(
response,
custom_message=
f"Image predict failed for the {title} model (ID: {model_id}).",
)
def test_search_by_image_url(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
response = stub.PostSearches(
service_pb2.PostSearchesRequest(
query=resources_pb2.Query(ands=[
resources_pb2.And(output=resources_pb2.Output(
input=resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL)))))
]),
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) > 0
assert input_.id in [hit.input.id for hit in response.hits]
def test_mixed_success_predict(channel):
stub = service_pb2_grpc.V2Stub(channel)
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL))),
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=NON_EXISTING_IMAGE_URL))),
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
assert response.status.code == status_code_pb2.MIXED_STATUS
assert response.outputs[0].status.code == status_code_pb2.SUCCESS
assert response.outputs[
1].status.code == status_code_pb2.INPUT_DOWNLOAD_FAILED
def test_patching_public_concept_fails(channel):
stub = service_pb2_grpc.V2Stub(channel)
patch_concepts_searches_response = stub.PatchConcepts(
service_pb2.PatchConceptsRequest(
action="overwrite",
concepts=[
resources_pb2.Concept(
id="ai_98Xb0K3q", # The ID of a public concept.
name="this new name won't be applied",
)
],
),
metadata=metadata(),
)
assert (patch_concepts_searches_response.status.code ==
status_code_pb2.StatusCode.CONN_DOES_NOT_EXIST)
assert patch_concepts_searches_response.status.description == "Resource does not exist"
def test_workflow_predict_image_url(channel):
stub = service_pb2_grpc.V2Stub(channel)
post_workflows_response = stub.PostWorkflowResults(
service_pb2.PostWorkflowResultsRequest(
workflow_id="General",
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(url=DOG_IMAGE_URL)))
],
output_config=resources_pb2.OutputConfig(max_concepts=3),
),
metadata=metadata(),
)
raise_on_failure(post_workflows_response)
assert len(
post_workflows_response.results[0].outputs[0].data.concepts) == 3
def test_predict_image_bytes(channel):
stub = service_pb2_grpc.V2Stub(channel)
with open(RED_TRUCK_IMAGE_FILE_PATH, "rb") as f:
file_bytes = f.read()
request = service_pb2.PostModelOutputsRequest(
model_id=GENERAL_MODEL_ID,
inputs=[
resources_pb2.Input(data=resources_pb2.Data(
image=resources_pb2.Image(base64=file_bytes)))
],
)
response = stub.PostModelOutputs(request, metadata=metadata())
raise_on_failure(response)
assert len(response.outputs[0].data.concepts) > 0
def test_search_by_predicted_concept_name(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
response = stub.PostSearches(
service_pb2.PostSearchesRequest(
query=resources_pb2.Query(ands=[
resources_pb2.And(output=resources_pb2.Output(
data=resources_pb2.Data(concepts=[
resources_pb2.Concept(name="dog", value=1)
])))
]),
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) > 0
assert input_.id in [hit.input.id for hit in response.hits]
def test_search_by_metadata(channel):
stub = service_pb2_grpc.V2Stub(channel)
with SetupImage(stub) as input_:
search_metadata = struct_pb2.Struct()
search_metadata.update({"another-key": {"inner-key": "inner-value"}})
response = stub.PostSearches(
service_pb2.PostSearchesRequest(
query=resources_pb2.Query(ands=[
resources_pb2.And(input=resources_pb2.Input(
data=resources_pb2.Data(metadata=search_metadata)))
]),
pagination=service_pb2.Pagination(page=1, per_page=1000),
),
metadata=metadata(),
)
raise_on_failure(response)
assert len(response.hits) > 0
assert input_.id in [hit.input.id for hit in response.hits]
