cnn-classification-service/cnn_classification_service/magic_doer.py

68 lines
1.8 KiB
Python

#!/usr/bin/env python3
import os
import logging
import tempfile
import requests
from cnn_classifier import Classifier
def run_everything(parameters: dict):
tag = parameters['tag']
sample_file_handle, sample_file_path = tempfile.mkstemp(prefix=f"{tag}_", suffix=".wav")
model_file_handle, model_file_path = tempfile.mkstemp(suffix=".json")
weights_file_handle, weights_file_path = tempfile.mkstemp(suffix=".h5")
try:
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}")
r = requests.get(f"http://model-service/model/cnn/$default")
r.raise_for_status()
with open(model_file_handle, 'wb') as f:
f.write(r.content)
r = requests.get(f"http://model-service/model/cnn/$default?weights")
r.raise_for_status()
with open(weights_file_handle, 'wb') as f:
f.write(r.content)
# magic happens here
classifier = Classifier(model_file_path, weights_file_path)
results = classifier.predict(sample_file_path)
finally: # bruuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuh
try:
os.remove(model_file_path)
except FileNotFoundError:
pass
try:
os.remove(weights_file_path)
except FileNotFoundError:
pass
try:
os.remove(sample_file_path)
except FileNotFoundError:
pass
response = {
"tag": tag,
"probability": 1.0 if results[0] == 'sturnus' else 0.0,
"model": ...
}
logging.info(f"Classification done!")
logging.debug(f"Results: {response}")
return response