def test_get_input_shape(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.get_input_shapes("xyc") == [[(1024, 512, 1)]]
assert model_session.get_input_shapes("cyx") == [[(1, 512, 1024)]]
assert model_session.get_input_shapes("c") == [[(1, )]]
assert model_session.get_input_shapes("tzyxc") == [[(1, 1, 512, 1024, 1)]]
def test_get_halo_returns_0_if_value_is_unspecified(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.get_halo(axes="xyz") == [256, 128, 0]
assert model_session.get_halo(axes="txyz") == [0, 256, 128, 0]
def test_get_halo_returns_values_specified_by_tags(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.get_halo(axes="xy") == [256, 128]
assert model_session.get_halo(axes="yx") == [128, 256]
assert model_session.get_halo(axes="yxc") == [128, 256, 1]
def test_known_classes(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.known_classes == list(range(1, 9))
def test_get_output_shape(pb_session):
"""shape = shape(input_tensor) * scale + 2 * offset"""
model_session = ModelSession(session=pb_session, factory=mock.Mock())
output_shape = model_session.get_output_shape()
assert output_shape == {"x": 1056, "y": 320, "c": 8}
def test_get_output_axes(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.output_axes == ["xyc"]
def test_has_training(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert not model_session.has_training
def test_num_classes(pb_session):
model_session = ModelSession(session=pb_session, factory=mock.Mock())
assert model_session.num_classes == 8