cnn-classification-service/cnn_classification_service/magic_doer.py

#!/usr/bin/env python3
import os
import logging
import tempfile
import requests
import time

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")
        response = None
        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
            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