def Predict(self, request: inference_pb2.PredictRequest,
context) -> inference_pb2.PredictResponse:
session = self._getModelSession(context, request.modelSessionId)
arr = converters.pb_tensor_to_numpy(request.tensor)
res = session.client.forward(arr)
pb_tensor = converters.numpy_to_pb_tensor(res)
return inference_pb2.PredictResponse(tensor=pb_tensor)
def _numpy_to_pb_tensor(arr):
"""
Makes sure that tensor was serialized/deserialized
"""
tensor = numpy_to_pb_tensor(arr)
parsed = inference_pb2.Tensor()
parsed.ParseFromString(tensor.SerializeToString())
return parsed
def predict(self, feature_image, roi, axistags=None):
"""
:param numpy.ndarray feature_image: classifier input
:param numpy.ndarray roi: ROI within feature_image
:param vigra.AxisTags axistags: axistags of feature_image
:return: probabilities
"""
assert isinstance(roi, numpy.ndarray)
logger.debug("predict tile shape: %s (axistags: %r)", feature_image.shape, axistags)
# translate roi axes todo: remove with tczyx standard
# output_axis_order = self._model_conf.output_axis_order
output_axis_order = self.output_axes
if "c" not in output_axis_order:
output_axis_order = "c" + output_axis_order
c_was_not_in_output_axis_order = True
else:
c_was_not_in_output_axis_order = False
roi = roi[:, [axistags.index(a) for a in output_axis_order]]
reordered_feature_image = reorder_axes(feature_image, from_axes_tags=axistags, to_axes_tags=self.input_axes)
try:
current_rq = Request._current_request()
resp = self.tiktorchClient.Predict.future(
inference_pb2.PredictRequest(
tensor=converters.numpy_to_pb_tensor(reordered_feature_image, axistags=self.input_axes),
modelSessionId=self.__session.id,
)
)
resp.add_done_callback(lambda o: current_rq._wake_up())
current_rq._suspend()
resp = resp.result()
result = converters.pb_tensor_to_numpy(resp.tensor)
except Exception:
logger.exception("Predict call failed")
return 0
logger.debug(f"Obtained a predicted block of shape {result.shape}")
if c_was_not_in_output_axis_order:
result = result[None, ...]
# make two channels out of single channel predictions
channel_axis = output_axis_order.find("c")
if result.shape[channel_axis] == 1:
result = numpy.concatenate((result, 1 - result), axis=channel_axis)
logger.debug(f"Changed shape of predicted block to {result.shape} by adding '1-p' channel")
shape_wo_halo = result.shape
result = result[roiToSlice(*roi)]
logger.debug(
f"Selected roi (start: {roi[0]}, stop: {roi[1]}) from result without halo {shape_wo_halo}. Now"
f" result has shape: ({result.shape})."
)
return reorder_axes(result, from_axes_tags=output_axis_order, to_axes_tags=axistags)
def test_call_predict(self, grpc_stub, pybio_dummy_model_bytes):
model = grpc_stub.CreateModelSession(valid_model_request(pybio_dummy_model_bytes))
arr = np.arange(32 * 32).reshape(1, 1, 32, 32)
expected = arr + 1
input_tensor = converters.numpy_to_pb_tensor(arr)
res = grpc_stub.Predict(inference_pb2.PredictRequest(modelSessionId=model.id, tensor=input_tensor))
grpc_stub.CloseModelSession(model)
assert_array_equal(expected, converters.pb_tensor_to_numpy(res.tensor))