def test_save_trained_model(self):
obj = {'model': 'knn', 'status': 'start'}
knn_model = KnnModel(spot_id=self.spot_id,
is_write=True,
knn_dir=self.knn_dir)
knn_model.save_trained_model(obj)
assert os.path.exists(self.knn_file_path)
def train_knn_model(self):
# load training data from csv file
csv_model = CsvModel(spot_id=self.spot_id,
is_write=False,
csv_dir=PATH_CSV_DIR)
classes, vector_images = csv_model.load_data()
# train
vector_images = np.array(vector_images)
neighbor = NearestNeighbors(n_neighbors=4)
logger.info({
'action': 'train_knn_model',
'status': 'start',
'message': 'start to train knn model'
})
neighbor.fit(vector_images)
logger.info({
'action': 'train_knn_model',
'status': 'end',
'message': 'finished to train knn model'
})
# save binary data as pickle
knn_model = KnnModel(spot_id=self.spot_id,
is_write=True,
knn_dir=PATH_KNN_DIR)
knn_model.save_trained_model(neighbor)
def test_load_trained_model(self):
obj = {'model': 'knn', 'status': 'start'}
knn_model = KnnModel(spot_id=self.spot_id,
is_write=False,
knn_dir=self.knn_dir)
knn_model.save_trained_model(obj)
actual_obj = knn_model.load_trained_model()
assert actual_obj == obj
def test_get_knn_file_from_s3(self, create_s3):
s3_key = os.path.join('knn', self.knn_file_name)
create_s3.Object(TestKnnModel.bucket, s3_key).put()
knn_model = KnnModel(spot_id=self.spot_id,
is_write=False,
knn_dir=self.knn_dir)
is_knn = knn_model.get_knn_file_from_s3()
assert is_knn is True
def inference_knn_model(self, data: list) -> list:
# preprocess vector data
data = np.array(data)
logger.info({'action': 'inference_knn_model', 'data': data})
data = data.reshape(1, -1)
logger.info({'action': 'inference_knn_model', 'data': data})
knn_model = KnnModel(spot_id=self.spot_id, is_write=False, knn_dir=PATH_KNN_DIR)
result_list = knn_model.inference(input_data=data)
# TODO: To be considered implementing function that exclude the prediction results by the point-to-point distance
result_classes = []
for result in result_list:
result_classes.append(result[0])
result_classes = list(set(result_classes))
logger.info({
'action': 'inference_knn_model',
'status': 'end',
'spot_id': self.spot_id,
'result_classes': result_classes
})
return result_classes
def test_get_knn_file_from_s3_create_file(self):
knn_model = KnnModel(spot_id=self.spot_id,
is_write=False,
knn_dir=self.knn_dir)
knn_model.get_knn_file_from_s3()
assert os.path.exists(self.knn_file_path)
def test_initialize_class_raises_no_directory(self):
with pytest.raises(NoDirectoryFoundError):
KnnModel(spot_id=1, is_write=False)
def test_initialize_class_raises(self):
with pytest.raises(NoSpotIdError):
KnnModel(spot_id=None, is_write=False, knn_dir='tmp')