forked from OlafenwaMoses/ImageAI
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_custom_classification.py
100 lines (84 loc) · 3.91 KB
/
test_custom_classification.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import os, sys
import cv2
from PIL import Image
import pytest
from os.path import dirname
sys.path.insert(1, os.path.join(dirname(dirname(os.path.abspath(__file__)))))
from imageai.Classification.Custom import CustomImageClassification
test_folder = dirname(os.path.abspath(__file__))
@pytest.mark.parametrize(
"image_input",
[
(os.path.join(test_folder, "data-images", "1.jpg")),
(cv2.imread(os.path.join(test_folder, "data-images", "1.jpg"))),
(Image.open(os.path.join(test_folder, "data-images", "1.jpg"))),
]
)
def test_recognition_model_mobilenetv2(image_input):
classifier = CustomImageClassification()
classifier.setModelTypeAsMobileNetV2()
classifier.setModelPath(os.path.join(test_folder, "data-models", "mobilenet_v2-idenprof-test_acc_0.85300_epoch-92.pt"))
classifier.setJsonPath(os.path.join(test_folder, "data-json", "idenprof_model_classes.json"))
classifier.loadModel()
predictions, probabilities = classifier.classifyImage(image_input=image_input, result_count=5)
assert isinstance(predictions, list)
assert isinstance(probabilities, list)
assert isinstance(predictions[0], str)
assert isinstance(probabilities[0], float)
@pytest.mark.parametrize(
"image_input",
[
(os.path.join(test_folder, "data-images", "1.jpg")),
(cv2.imread(os.path.join(test_folder, "data-images", "1.jpg"))),
(Image.open(os.path.join(test_folder, "data-images", "1.jpg"))),
]
)
def test_recognition_model_resnet(image_input):
classifier = CustomImageClassification()
classifier.setModelTypeAsResNet50()
classifier.setModelPath(os.path.join(test_folder, "data-models", "resnet50-idenprof-test_acc_0.78200_epoch-91.pt"))
classifier.setJsonPath(os.path.join(test_folder, "data-json", "idenprof_model_classes.json"))
classifier.loadModel()
predictions, probabilities = classifier.classifyImage(image_input=image_input, result_count=5)
assert isinstance(predictions, list)
assert isinstance(probabilities, list)
assert isinstance(predictions[0], str)
assert isinstance(probabilities[0], float)
@pytest.mark.parametrize(
"image_input",
[
(os.path.join(test_folder, "data-images", "1.jpg")),
(cv2.imread(os.path.join(test_folder, "data-images", "1.jpg"))),
(Image.open(os.path.join(test_folder, "data-images", "1.jpg"))),
]
)
def test_recognition_model_inceptionv3(image_input):
classifier = CustomImageClassification()
classifier.setModelTypeAsInceptionV3()
classifier.setModelPath(os.path.join(test_folder, "data-models", "inception_v3-idenprof-test_acc_0.81050_epoch-92.pt"))
classifier.setJsonPath(os.path.join(test_folder, "data-json", "idenprof_model_classes.json"))
classifier.loadModel()
predictions, probabilities = classifier.classifyImage(image_input=image_input, result_count=5)
assert isinstance(predictions, list)
assert isinstance(probabilities, list)
assert isinstance(predictions[0], str)
assert isinstance(probabilities[0], float)
@pytest.mark.parametrize(
"image_input",
[
(os.path.join(test_folder, "data-images", "1.jpg")),
(cv2.imread(os.path.join(test_folder, "data-images", "1.jpg"))),
(Image.open(os.path.join(test_folder, "data-images", "1.jpg"))),
]
)
def test_recognition_model_densenet(image_input):
classifier = CustomImageClassification()
classifier.setModelTypeAsDenseNet121()
classifier.setModelPath(os.path.join(test_folder, "data-models", "densenet121-idenprof-test_acc_0.82550_epoch-95.pt"))
classifier.setJsonPath(os.path.join(test_folder, "data-json", "idenprof_model_classes.json"))
classifier.loadModel()
predictions, probabilities = classifier.classifyImage(image_input=image_input, result_count=5)
assert isinstance(predictions, list)
assert isinstance(probabilities, list)
assert isinstance(predictions[0], str)
assert isinstance(probabilities[0], float)