frederik
223 lines
6.4 KiB
Python
223 lines
6.4 KiB
Python
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'
|