cnn-classification-service/cnn_classification_service/magic_doer.py

#!/usr/bin/env python3
import os
import logging
import tempfile
from requests_opentracing import SessionTracing
import opentracing  # ez kell ide hogy a session tracer jolegyen
import requests
import time
from urllib.parse import urljoin

from config import Config

from classifier_cache import ClassifierCache


class MagicDoer:
    classifier_cache = ClassifierCache(Config.MODEL_INFO_URL)
    requests_session = SessionTracing(propagate=True)

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

            with opentracing.tracer.start_span('downloadSample', child_of=span):
                # Download Sample
                object_path = urljoin(Config.STORAGE_SERVICE_URL, f"object/{tag}")

                logging.info(f"Downloading sample: {tag} from {object_path}")
                r = cls.requests_session.get(object_path)
                with open(sample_file_handle, 'wb') as f:
                    f.write(r.content)

                logging.debug(f"Downloaded sample to {sample_file_path}")

            with opentracing.tracer.start_span('loadClassifier', child_of=span):
                # Get a classifier that uses the default model
                model_details, classifier = cls.classifier_cache.get_default_classifier()

            with opentracing.tracer.start_span('runClassifier', child_of=span) as child_span:
                # do the majic
                classification_start_time = time.time()
                predicted_class_name, labeled_predictions = classifier.predict(sample_file_path, child_span)
                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

        span.log_kv({'event': 'sampleFileDeleted', 'sampleTag': tag})

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

        return response