def test_efficientdet():
model: EfficientDetector = try_cuda(EfficientDetector(n_classes=n_classes, score_threshold=1e-2))
test_image = gen_random_tensor(1, image_size)
result = model.predict_bboxes(test_image)
assert result[0].ndim == 2 and result[0].shape[
-1] == 6, "Detection model result must contain (xmin, ymin, xmax, ymax, class, score)"
def _test_classification(model: ClassificationModel, batch_size: int, n_classes=11, image_size=(96, 96)):
expected_prob_tensor_shape = (batch_size, n_classes)
expected_label_tensor_shape = (batch_size,)
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs = model.predict_label(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape, "output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")
def _test_attention_classification(model: AttentionClassificationModel, batch_size: int, n_classes=11,
image_size=(96, 96)):
model = try_cuda(model)
expected_prob_tensor_shape = (batch_size, n_classes)
expected_label_tensor_shape = (batch_size,)
expected_attention_shape_length = 3 # (batch size, height, width)
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs, attention_map = model.predict_label_and_heatmap(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape, "output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")
assert len(attention_map.shape) == expected_attention_shape_length
def _test_segmentation(model: SegmentationModel,
batch_size: int,
n_classes=11,
image_size=(256, 256)):
model = try_cuda(model)
expected_prob_tensor_shape = (batch_size, n_classes) + image_size
expected_label_tensor_shape = (batch_size, ) + image_size
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs = model.predict_index_image(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape,
"output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")