def test_predict(self, mock_load, mock_predict, mock_preprocess_ds):
mock_load.return_value = ParallelPostFit()
mock_preprocess_ds.return_value = [1, 2, 3]
testable = SKLinearImageModel(pkl_file='trained_models/hog_sklearn.pkl')
testable.predict(self.X)
mock_load.assert_called_once()
mock_predict.assert_called_once_with([1, 2, 3])
def test_train(self, mock_preprocess_ds, mock_fit, mock_predict_proba):
X_train, X_test, y_train, y_test = train_test_split(
self.X,
self.y,
test_size=0.2,
shuffle=True,
random_state=42,
)
X_train_prepared = X_train * 0.2
mock_preprocess_ds.return_value = X_train_prepared
mock_fit.return_value = ParallelPostFit()
testable = SKLinearImageModel(pkl_file=None)
testable.train(X_train, y_train, X_test, y_test, verbose=False)
np.testing.assert_array_equal(X_train_prepared,
mock_fit.call_args[0][0])
np.testing.assert_array_equal(y_train,
mock_fit.call_args[0][1])
mock_predict_proba.assert_called_once()
epilog=textwrap.dedent(DOC))
arg_parser.add_argument('--dir',
help='Path to the input directory with image files to train the model on.')
arg_parser.add_argument('--pkl_file', default='trained_models/fitted_model.pkl',
help='Pickle file name to save the trained model')
arg_parser.add_argument('--batch_size', default=1000,
help='Batch size for large amount of training data')
parsed_args = arg_parser.parse_args(sys.argv[1:])
input_directory = parsed_args.dir
pkl_file = parsed_args.pkl_file
batch_size = parsed_args.batch_size
try:
dir_parser = DirectoryParser(input_directory)
except ValueError as e:
print(*e.args)
sys.exit(1)
print(dir_parser)
model = SKLinearImageModel(pkl_file=None)
feature_extractor = ImageFeatureExtractor()
batches = feature_extractor.transform_image_to_dataset(dir_parser.full_path_image_files,
batch_size=batch_size)
X_train, y_train, X_test, y_test = feature_extractor.combine_batches(batches=batches)
model.train(X_train, y_train, X_test, y_test)
model.save(pkl_file)
help='Limit to accept probability to predict a class.')
arg_parser.add_argument('dir',
help='Path to the input directory with image files.')
parsed_args = arg_parser.parse_args(sys.argv[1:])
input_directory = parsed_args.dir
pkl_file = parsed_args.pkl_file
probability_threshold = float(parsed_args.probability_threshold)
try:
dir_parser = DirectoryParser(input_directory)
except ValueError as e:
print(*e.args)
sys.exit(1)
print(f"Data dir: {dir_parser}")
feature_extractor = ImageFeatureExtractor()
if parsed_args.model_type == ModelUsageType.sk:
model = SKLinearImageModel(pkl_file=pkl_file,
probability_threshold=probability_threshold)
X, _ = feature_extractor.transform_image_to_dataset(dir_parser.full_path_image_files)[0]
else:
model = TFModel(pre_train=pkl_file,
probability_threshold=probability_threshold)
X, _ = feature_extractor.transform_image_to_dataset(dir_parser.full_path_image_files,
image_size=model.image_shape[:-1])[0]
prediction = model.predict(X)
for filepath, prediction in zip(dir_parser.full_path_image_files, prediction):
print(f"{filepath}: {prediction.value}")
def test_predict_proba_to_label_unknown(self):
testable = SKLinearImageModel(pkl_file=None)
proba = (0.45, 0.65)
self.assertEqual('unknown', testable._predict_proba_to_label(proba))
def test_predict_proba_to_label_dog(self):
testable = SKLinearImageModel(pkl_file=None)
proba = (0.15, 0.85)
self.assertEqual('dog', testable._predict_proba_to_label(proba))
def test_load(self, mock_joblib):
testable = SKLinearImageModel(pkl_file=None)
testable.load('1.pkl')
mock_joblib.assert_called_once_with('1.pkl')
def test_save(self, mock_joblib):
testable = SKLinearImageModel(pkl_file=None)
testable.classifier = mock.MagicMock()
testable.save('1.pkl')
mock_joblib.assert_called_once_with(testable.classifier, '1.pkl')