84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import logging
|
|
import tempfile
|
|
import requests
|
|
|
|
from urllib.parse import urljoin
|
|
|
|
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:
|
|
|
|
# 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}")
|
|
|
|
# Download model
|
|
|
|
model_root_url = "http://model-service/model/cnn/$default"
|
|
|
|
logging.debug("Fetching model info...")
|
|
r = requests.get(model_root_url)
|
|
r.raise_for_status()
|
|
|
|
model_details = r.json()
|
|
|
|
logging.debug("Fetching model file...")
|
|
r = requests.get(urljoin(model_root_url, model_details['files']['model'])) # Fun fact: this would support external urls
|
|
r.raise_for_status()
|
|
|
|
with open(model_file_handle, 'wb') as f:
|
|
f.write(r.content)
|
|
|
|
logging.debug("Fetching weights file...")
|
|
r = requests.get(urljoin(model_root_url, model_details['files']['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] == model_details['target_class_name'] else 0.0,
|
|
"model": model_details['id']
|
|
}
|
|
|
|
logging.info(f"Classification done!")
|
|
logging.debug(f"Results: {response}")
|
|
|
|
return response
|