cnn-classification-service/cnn_classification_service/magic_doer.py

#!/usr/bin/env python3
import os
import logging
import tempfile
import requests
import time
from urllib.parse import urljoin

from classifier_cache import ClassifierCache


class MagicDoer:
    classifier_cache = ClassifierCache(
        os.environ.get("MODEL_INFO_URL", "http://model-service/model/cnn/$default")
    )

    @classmethod
    def run_everything(cls, parameters: dict) -> dict:
        tag = parameters['tag']
        sample_file_handle, sample_file_path = tempfile.mkstemp(prefix=f"{tag}_", suffix=".wav", dir="/dev/shm")
        response = None
        try:

            # Download Sample
            storage_service_url = os.environ.get("STORAGE_SERVICE_URL", "http://storage-service/")
            object_path = urljoin(storage_service_url, f"object/{tag}")

            logging.info(f"Downloading sample: {tag} from {object_path}")
            r = requests.get(object_path)
            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
            classification_start_time = time.time()
            predicted_class_name, labeled_predictions = classifier.predict(sample_file_path)
            classification_duration = time.time() - classification_start_time

            response = {
                "tag": tag,
                "probability": labeled_predictions[model_details['target_class_name']],
                "all_predictions": labeled_predictions,
                "class": predicted_class_name,
                "model": model_details['id'],
                "classification_duration": classification_duration
            }

        finally:
            try:
                os.remove(sample_file_path)
            except FileNotFoundError:
                pass

        if not response:
            logging.error("Something went wrong during classification!")
        else:
            logging.info(f"Classification done!")
            logging.debug(f"Results: {response}")

        return response