diff --git a/scripts/utils/models.py b/scripts/utils/models.py new file mode 100644 index 0000000..45836b1 --- /dev/null +++ b/scripts/utils/models.py @@ -0,0 +1,222 @@ +import logging +import math +import operator +import os + +import numpy as np + +from .helpers import get_settings, get_model_labels, MODEL_PATH + +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' +os.environ['CUDA_VISIBLE_DEVICES'] = '' +np.set_printoptions(legacy="1.21") + +try: + import tflite_runtime.interpreter as tflite +except ImportError: + from tensorflow import lite as tflite + +log = logging.getLogger(__name__) + + +def get_model(model=None): + conf = get_settings() + if model is None: + model = conf['MODEL'] + + if model == 'BirdNET_6K_GLOBAL_MODEL': + return BirdNetV1(conf.getfloat('SENSITIVITY')) + elif model == 'BirdNET_GLOBAL_6K_V2.4_Model_FP16': + return BirdNetV2_4(conf.getfloat('SENSITIVITY')) + + +def get_meta_model(model=None, version=None): + conf = get_settings() + if model is None: + model = conf['MODEL'] + if version is None: + version = conf.getint('DATA_MODEL_VERSION') + + if model != 'BirdNET_GLOBAL_6K_V2.4_Model_FP16': + return None + + if version == 1: + return MDataModel1(conf.getfloat('SF_THRESH')) + elif version == 2: + return MDataModel2(conf.getfloat('SF_THRESH')) + + +class Basemodel: + chunk_duration = None + sample_rate = None + model_name = None + _input_layer = 0 + _output_layer = 0 + + def __init__(self): + model_path = os.path.join(MODEL_PATH, f'{self.model_name}.tflite') + self.interpreter = tflite.Interpreter(model_path) + self.interpreter.allocate_tensors() + input_details = self.interpreter.get_input_details() + output_details = self.interpreter.get_output_details() + + self._input_layer_idx = input_details[self._input_layer]['index'] + self._output_layer_idx = output_details[self._output_layer]['index'] + + self.labels = get_model_labels(self.model_name) + + def label(self, logits): + p_labels = dict(zip(self.labels, logits)) + return sorted(p_labels.items(), key=operator.itemgetter(1), reverse=True) + + def predict(self, chunk): + raise NotImplementedError + + def set_meta_data(self, lat, lon, week): + pass + + def get_species_list(self): + return [] + + +class BirdNet(Basemodel): + chunk_duration = 3 + sample_rate = 48000 + + def __init__(self, sens): + super().__init__() + + self._mdata_model = self._set_meta_model() + + self._sensitivity = max(0.5, min(1.0 - (sens - 1.0), 1.5)) + + def scale(self, logits): + return 1 / (1.0 + np.exp(-self._sensitivity * logits)) + + def _set_meta_model(self): + return None + + +class BirdNetV1(BirdNet): + model_name = 'BirdNET_6K_GLOBAL_MODEL' + + def __init__(self, sens): + super().__init__(sens) + self._mdata = None + self._mdata_params = None + + def _set_meta_model(self): + input_details = self.interpreter.get_input_details() + return input_details[1]['index'] + + def predict(self, chunk): + self.interpreter.set_tensor(self._input_layer_idx, np.array(chunk, dtype='float32')[np.newaxis, :]) + self.interpreter.set_tensor(self._mdata_model, np.array(self._mdata, dtype='float32')) + + self.interpreter.invoke() + logits = self.interpreter.get_tensor(self._output_layer_idx)[0] + + return self.label(self.scale(logits)) + + def _convert_metadata(self, m): + # Convert week to cosine + if 1 <= m[2] <= 48: + m[2] = math.cos(math.radians(m[2] * 7.5)) + 1 + else: + m[2] = -1 + + # Add binary mask + mask = np.ones((3,)) + if m[0] == -1 or m[1] == -1: + mask = np.zeros((3,)) + if m[2] == -1: + mask[2] = 0.0 + + return np.concatenate([m, mask]) + + def set_meta_data(self, lat, lon, week): + if self._mdata_params != [lat, lon, week]: + self._mdata_params = [lat, lon, week] + # Convert and prepare metadata + mdata = self._convert_metadata(np.array([lat, lon, week])) + self._mdata = np.expand_dims(mdata, 0) + + +class BirdNetV2_4(BirdNet): + model_name = 'BirdNET_GLOBAL_6K_V2.4_Model_FP16' + + def _set_meta_model(self): + return get_meta_model() + + def predict(self, chunk): + self.interpreter.set_tensor(self._input_layer_idx, np.array(chunk, dtype='float32')[np.newaxis, :]) + + self.interpreter.invoke() + logits = self.interpreter.get_tensor(self._output_layer_idx)[0] + + return self.label(self.scale(logits)) + + def set_meta_data(self, lat, lon, week): + self._mdata_model.set_meta_data(lat, lon, week) + + def get_species_list(self): + return self._mdata_model.get_species_list(self.labels) + + +class MDataModel: + model_name = None + + def __init__(self, sf_thresh): + model_path = os.path.join(MODEL_PATH, f'{self.model_name}.tflite') + self.interpreter = tflite.Interpreter(model_path) + self.interpreter.allocate_tensors() + input_details = self.interpreter.get_input_details() + output_details = self.interpreter.get_output_details() + + self._input_layer_idx = input_details[0]['index'] + self._output_layer_idx = output_details[0]['index'] + self._sf_thresh = sf_thresh + + self._mdata_params = None + self._mdata = None + + def set_meta_data(self, lat, lon, week): + if self._mdata_params != (lat, lon, week): + self._mdata = None + self._mdata_params = (lat, lon, week) + + def get_species_list_details(self, labels): + if self._mdata is None: + lat, lon, week = self._mdata_params + sample = np.expand_dims(np.array([lat, lon, week], dtype='float32'), 0) + + # Run inference + self.interpreter.set_tensor(self._input_layer_idx, sample) + self.interpreter.invoke() + + l_filter = self.interpreter.get_tensor(self._output_layer_idx)[0] + + # Apply threshold + l_filter = np.where(l_filter >= float(self._sf_thresh), l_filter, 0) + + # Zip with labels + l_filter = list(zip(l_filter, labels)) + + # Sort by filter value + l_filter = sorted(l_filter, key=lambda x: x[0], reverse=True) + + self._mdata = [s for s in l_filter if s[0] >= self._sf_thresh] + + return self._mdata + + def get_species_list(self, labels): + l_filter = self.get_species_list_details(labels) + return [s[1].split('_')[0] for s in l_filter] + + +class MDataModel1(MDataModel): + model_name = 'BirdNET_GLOBAL_6K_V2.4_MData_Model_FP16' + + +class MDataModel2(MDataModel): + model_name = 'BirdNET_GLOBAL_6K_V2.4_MData_Model_V2_FP16'