cnn-classification-service/cnn_classification_service/magic_doer.py
marcsello 751fbdc4a5
All checks were successful
continuous-integration/drone/push Build is passing
Implemented classifier cache
2020-10-06 00:41:54 +02:00

50 lines
1.4 KiB
Python

#!/usr/bin/env python3
import os
import logging
import tempfile
import requests
from classifier_cache import ClassifierCache
class MagicDoer:
classifier_cache = ClassifierCache()
@classmethod
def run_everything(cls, parameters: dict) -> dict:
tag = parameters['tag']
sample_file_handle, sample_file_path = tempfile.mkstemp(prefix=f"{tag}_", suffix=".wav")
try:
# Download Sample
logging.info(f"Downloading sample: {tag}")
r = requests.get(f"http://storage-service/object/{tag}")
with open(sample_file_handle, 'wb') as f:
f.write(r.content)
logging.debug(f"Downloaded sample to {sample_file_path}")
# Get a classifier that uses the default model
model_details, classifier = cls.classifier_cache.get_default_classifier()
# do the majic
results = classifier.predict(sample_file_path)
finally:
try:
os.remove(sample_file_path)
except FileNotFoundError:
pass
response = {
"tag": tag,
"probability": 1.0 if results[0] == model_details['target_class_name'] else 0.0,
"model": model_details['id']
}
logging.info(f"Classification done!")
logging.debug(f"Results: {response}")
return response