From 24bda466bdae6283df5b4b95705601e8a4130650 Mon Sep 17 00:00:00 2001 From: Jason Parham Date: Wed, 27 Mar 2024 11:53:43 -0700 Subject: [PATCH 1/8] wip --- .gitattributes | 1 + .gitignore | 5 + assets/model.mobilenet.onnx | 3 + bats_ai/core/admin/spectrogram.py | 2 + .../migrations/0009_spectrogram_colormap.py | 17 ++ bats_ai/core/models/__init__.py | 3 +- bats_ai/core/models/recording.py | 43 +++- bats_ai/core/models/spectrogram.py | 187 +++++++++++++++--- bats_ai/core/rest/__init__.py | 11 +- bats_ai/core/rest/spectrogram.py | 21 ++ bats_ai/core/tasks.py | 16 +- bats_ai/core/views/recording.py | 24 ++- bats_ai/settings.py | 1 + bats_ai/urls.py | 13 +- client/yarn.lock | 11 +- dev/django.Dockerfile | 2 + setup.py | 7 + 17 files changed, 315 insertions(+), 52 deletions(-) create mode 100644 assets/model.mobilenet.onnx create mode 100644 bats_ai/core/migrations/0009_spectrogram_colormap.py create mode 100644 bats_ai/core/rest/spectrogram.py diff --git a/.gitattributes b/.gitattributes index 84fe760..f8770d0 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ assets/example.wav filter=lfs diff=lfs merge=lfs -text +assets/model.mobilenet.onnx filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index 78d3c39..58607bd 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,8 @@ /**/*.shp /**/*.shx /**/*.csv +models/datasets/ +models/spectrograms/ +models/ignore/ +models/*.jpg +models/*.pkl diff --git a/assets/model.mobilenet.onnx b/assets/model.mobilenet.onnx new file mode 100644 index 0000000..898cca5 --- /dev/null +++ b/assets/model.mobilenet.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6e36494aa6457f8d93bb882e90fe7bf2100484fdaa2085d15856217ba1fcbf1 +size 12135629 diff --git a/bats_ai/core/admin/spectrogram.py b/bats_ai/core/admin/spectrogram.py index 2586ebb..6f2f601 100644 --- a/bats_ai/core/admin/spectrogram.py +++ b/bats_ai/core/admin/spectrogram.py @@ -8,6 +8,7 @@ class SpectrogramAdmin(admin.ModelAdmin): list_display = [ 'pk', 'recording', + 'colormap', 'image_file', 'width', 'height', @@ -20,6 +21,7 @@ class SpectrogramAdmin(admin.ModelAdmin): autocomplete_fields = ['recording'] readonly_fields = [ 'recording', + 'colormap', 'image_file', 'created', 'modified', diff --git a/bats_ai/core/migrations/0009_spectrogram_colormap.py b/bats_ai/core/migrations/0009_spectrogram_colormap.py new file mode 100644 index 0000000..a89259e --- /dev/null +++ b/bats_ai/core/migrations/0009_spectrogram_colormap.py @@ -0,0 +1,17 @@ +# Generated by Django 4.1.13 on 2024-03-12 21:53 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ('core', '0008_grtscells_recording_recorded_time'), + ] + + operations = [ + migrations.AddField( + model_name='spectrogram', + name='colormap', + field=models.CharField(max_length=20, null=True), + ), + ] diff --git a/bats_ai/core/models/__init__.py b/bats_ai/core/models/__init__.py index 553e7cb..5979660 100644 --- a/bats_ai/core/models/__init__.py +++ b/bats_ai/core/models/__init__.py @@ -1,7 +1,7 @@ from .annotations import Annotations from .grts_cells import GRTSCells from .image import Image -from .recording import Recording +from .recording import Recording, colormap from .recording_annotation_status import RecordingAnnotationStatus from .species import Species from .spectrogram import Spectrogram @@ -16,4 +16,5 @@ 'Spectrogram', 'TemporalAnnotations', 'GRTSCells', + 'colormap', ] diff --git a/bats_ai/core/models/recording.py b/bats_ai/core/models/recording.py index 19cf559..7ac791a 100644 --- a/bats_ai/core/models/recording.py +++ b/bats_ai/core/models/recording.py @@ -1,7 +1,46 @@ +import logging + from django.contrib.auth.models import User from django.contrib.gis.db import models from django_extensions.db.models import TimeStampedModel +logger = logging.getLogger(__name__) + + +COLORMAP = None +COLORMAP_ALLOWED = [None, 'gist_yarg', 'gnuplot2'] + + +class colormap: + def __init__(self, colormap=None): + # Helpful aliases + if colormap in ['none', 'default', 'dark']: + colormap = None + if colormap in ['light']: + colormap = 'gist_yarg' + if colormap in ['heatmap']: + colormap = 'turbo' + + # Supported colormaps + if colormap not in COLORMAP_ALLOWED: + logging.warning(f'Substituted requested {colormap} colormap to default') + logging.warning('See COLORMAP_ALLOWED') + colormap = None + + self.colormap = colormap + self.previous = None + + def __enter__(self): + global COLORMAP + + self.previous = COLORMAP + COLORMAP = self.colormap + + def __exit__(self, exc_type, exc_value, exc_tb): + global COLORMAP + + COLORMAP = self.previous + # TimeStampedModel also provides "created" and "modified" fields class Recording(TimeStampedModel, models.Model): @@ -25,7 +64,7 @@ def has_spectrogram(self): def spectrograms(self): from bats_ai.core.models import Spectrogram - query = Spectrogram.objects.filter(recording=self).order_by('-created') + query = Spectrogram.objects.filter(recording=self, colormap=COLORMAP).order_by('-created') return query.all() @property @@ -35,7 +74,7 @@ def spectrogram(self): spectrograms = self.spectrograms if len(spectrograms) == 0: - Spectrogram.generate(self) + Spectrogram.generate(self, colormap=COLORMAP) spectrograms = self.spectrograms assert len(spectrograms) == 1 diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index 9ee4ab6..c6892e7 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -12,8 +12,9 @@ import matplotlib.pyplot as plt import numpy as np import scipy +import tqdm -from bats_ai.core.models import Recording +from bats_ai.core.models import Annotations, Recording FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -29,9 +30,13 @@ class Spectrogram(TimeStampedModel, models.Model): duration = models.IntegerField() # milliseconds frequency_min = models.IntegerField() # hz frequency_max = models.IntegerField() # hz + colormap = models.CharField(max_length=20, blank=False, null=True) @classmethod - def generate(cls, recording): + def generate(cls, recording, colormap=None, dpi=600): + """ + Ref: https://matplotlib.org/stable/users/explain/colors/colormaps.html + """ wav_file = recording.audio_file try: if isinstance(wav_file, FieldFile): @@ -45,13 +50,17 @@ def generate(cls, recording): print(e) return None + import IPython + + IPython.embed() + # Function to take a signal and return a spectrogram. size = int(0.001 * sr) # 1.0ms resolution - size = 2 ** math.ceil(math.log(size, 2)) + size = 2 ** (math.ceil(math.log(size, 2)) + 0) hop_length = int(size / 4) # Short-time Fourier Transform - window = librosa.stft(sig, n_fft=size, window='hamming') + window = librosa.stft(sig, n_fft=size, hop_length=hop_length, window='hamming') # Calculating and processing data for the spectrogram. window = np.abs(window) ** 2 @@ -79,14 +88,13 @@ def generate(cls, recording): freq_high = int(FREQ_MAX + FREQ_PAD) vmin = window.min() vmax = window.max() - dpi = 300 - chunksize = int(5e3) + chunksize = int(2e3) arange = np.arange(chunksize, window.shape[1], chunksize) chunks = np.array_split(window, arange, axis=1) imgs = [] - for chunk in chunks: + for chunk in tqdm.tqdm(chunks): h, w = chunk.shape alpha = 3 figsize = (int(math.ceil(w / h)) * alpha + 1, alpha) @@ -94,17 +102,22 @@ def generate(cls, recording): ax = plt.axes() plt.margins(0) + kwargs = { + 'sr': sr, + 'n_fft': size, + 'hop_length': hop_length, + 'x_axis': 's', + 'y_axis': 'fft', + 'ax': ax, + 'vmin': vmin, + 'vmax': vmax, + } + # Plot - librosa.display.specshow( - chunk, - sr=sr, - hop_length=hop_length, - x_axis='s', - y_axis='linear', - ax=ax, - vmin=vmin, - vmax=vmax, - ) + if colormap is None: + librosa.display.specshow(chunk, **kwargs) + else: + librosa.display.specshow(chunk, cmap=colormap, **kwargs) ax.set_ylim(freq_low, freq_high) ax.axis('off') @@ -132,12 +145,44 @@ def generate(cls, recording): w, h = img.size # ratio = dpi / h # w_ = int(round(w * ratio)) - w_ = int(4.0 * duration * 1e3) - h_ = int(dpi) + w_ = int(4.0 * duration * 1e3) * 2 + h_ = int(dpi) * 2 img = img.resize((w_, h_), resample=Image.Resampling.LANCZOS) w, h = img.size - # img.save('temp.jpg') + img.save(f'temp.{colormap}.png') + + img_filtered = cv2.imread(f'temp.{colormap}.jpg', cv2.IMREAD_GRAYSCALE) + img_filtered = img_filtered.astype(np.float32) + assert img_filtered.min() == 0 + assert img_filtered.max() == 255 + img_filtered = img_filtered.astype(np.float32) + img_filtered /= 0.9 + img_filtered[img_filtered < 0] = 0 + img_filtered[img_filtered > 255] = 255 + img_filtered = 255.0 - img_filtered + + kernel = np.ones((3, 3), np.uint8) + # img_filtered = cv2.morphologyEx(img_filtered, cv2.MORPH_OPEN, kernel) + img_filtered = cv2.erode(img_filtered, kernel, iterations=1) + + img_filtered = np.sqrt(img_filtered / 255.0) * 255.0 + + img_filtered = np.around(img_filtered).astype(np.uint8) + # img_filtered = cv2.resize(img_filtered, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA) + # img_filtered = cv2.medianBlur(img_filtered, 3) + # img_filtered = cv2.resize(img_filtered, None, fx=4.0, fy=4.0, interpolation=cv2.INTER_LINEAR) + + mask = img_filtered < 255 * 0.1 + img_filtered = cv2.applyColorMap(img_filtered, cv2.COLORMAP_TURBO) + img_filtered[mask] = [0, 0, 0] + + cv2.imwrite('temp.png', img_filtered) + + # img_filtered = img.resize((w_, h_), resample=Image.Resampling.LANCZOS) + # img_filtered = img_filtered.filter(ImageFilter.MedianFilter(size=3)) + # img_filtered = img_filtered.resize((w, h), resample=Image.Resampling.BILINEAR) + # img_filtered.save(f'temp.{colormap}.median.jpg') buf = io.BytesIO() img.save(buf, format='JPEG', quality=80) @@ -154,6 +199,7 @@ def generate(cls, recording): duration=math.ceil(duration * 1e3), frequency_min=freq_low, frequency_max=freq_high, + colormap=colormap, ) spectrogram.save() @@ -161,11 +207,17 @@ def generate(cls, recording): def compressed(self): img = self.image_np + annotations = Annotations.objects.filter(recording=self.recording) + threshold = 0.5 while True: canvas = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) canvas = canvas.astype(np.float32) + is_light = np.median(canvas) > 128.0 + if is_light: + canvas = 255.0 - canvas + amplitude = canvas.max(axis=0) amplitude -= amplitude.min() amplitude /= amplitude.max() @@ -179,10 +231,20 @@ def compressed(self): canvas *= amplitude canvas = np.around(canvas).astype(np.uint8) + width = canvas.shape[1] + for annotation in annotations: + start = annotation.start_time / self.duration + stop = annotation.end_time / self.duration + + start = int(np.around(start * width)) + stop = int(np.around(stop * width)) + canvas[:, start : stop + 1] = 255.0 + mask = canvas.max(axis=0) mask = scipy.signal.medfilt(mask, 3) mask[0] = 0 mask[-1] = 0 + starts = [] stops = [] for index in range(1, len(mask) - 1): @@ -268,14 +330,20 @@ def compressed(self): canvas = Image.fromarray(canvas, 'RGB') - # canvas.save('temp.jpg') + canvas.save('temp.compressed.jpg') buf = io.BytesIO() canvas.save(buf, format='JPEG', quality=80) buf.seek(0) img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8') - return img_base64, starts_, stops_, widths, total_width + metadata = { + 'starts': starts_, + 'stops': stops_, + 'widths': widths, + 'length': total_width, + } + return canvas, img_base64, metadata @property def image_np(self): @@ -296,3 +364,78 @@ def base64(self): img_base64 = base64.b64encode(img).decode('utf-8') return img_base64 + + def predict(self): + import json + import os + + import onnx + import onnxruntime as ort + import tqdm + + img, _, _ = self.compressed + + relative = ('..',) * 4 + asset_path = os.path.abspath(os.path.join(__file__, *relative, 'assets')) + + onnx_filename = os.path.join(asset_path, 'model.mobilenet.onnx') + assert os.path.exists(onnx_filename) + + session = ort.InferenceSession( + onnx_filename, + providers=[ + ( + 'CUDAExecutionProvider', + { + 'cudnn_conv_use_max_workspace': '1', + 'device_id': 0, + 'cudnn_conv_algo_search': 'HEURISTIC', + }, + ), + 'CPUExecutionProvider', + ], + ) + + img = np.array(img) + + h, w, c = img.shape + ratio_y = 224 / h + ratio_x = ratio_y * 0.5 + raw = cv2.resize(img, None, fx=ratio_x, fy=ratio_y, interpolation=cv2.INTER_LANCZOS4) + + h, w, c = raw.shape + if w <= h: + canvas = np.zeros((h, h + 1, 3), dtype=raw.dtype) + canvas[:, :w, :] = raw + raw = canvas + h, w, c = raw.shape + + inputs_ = [] + for index in range(0, w - h, 100): + inputs_.append(raw[:, index : index + h, :]) + inputs_.append(raw[:, -h:, :]) + inputs_ = np.array(inputs_) + + chunksize = 1 + chunks = np.array_split(inputs_, np.arange(chunksize, len(inputs_), chunksize)) + outputs = [] + for chunk in tqdm.tqdm(chunks, desc='Inference'): + outputs_ = session.run( + None, + {'input': chunk}, + ) + outputs.append(outputs_[0]) + outputs = np.vstack(outputs) + outputs = outputs.mean(axis=0) + + model = onnx.load(onnx_filename) + mapping = json.loads(model.metadata_props[0].value) + labels = [mapping['forward'][str(index)] for index in range(len(mapping['forward']))] + + prediction = np.argmax(outputs) + label = labels[prediction] + score = outputs[prediction] + + confs = dict(zip(labels, outputs)) + + return label, score, confs diff --git a/bats_ai/core/rest/__init__.py b/bats_ai/core/rest/__init__.py index a3bdd5b..f94b8a5 100644 --- a/bats_ai/core/rest/__init__.py +++ b/bats_ai/core/rest/__init__.py @@ -1,5 +1,10 @@ +from rest_framework import routers + from .image import ImageViewSet +from .spectrogram import SpectrogramViewSet + +__all__ = ['ImageViewSet', 'SpectrogramViewSet'] -__all__ = [ - 'ImageViewSet', -] +rest = routers.SimpleRouter() +# rest.register(r'images', ImageViewSet) +rest.register(r'spectrograms', SpectrogramViewSet) diff --git a/bats_ai/core/rest/spectrogram.py b/bats_ai/core/rest/spectrogram.py new file mode 100644 index 0000000..a9bf136 --- /dev/null +++ b/bats_ai/core/rest/spectrogram.py @@ -0,0 +1,21 @@ +from django_large_image.rest import LargeImageFileDetailMixin +from rest_framework import mixins, serializers, viewsets + +from bats_ai.core.models import Spectrogram + + +class SpectrogramSerializer(serializers.ModelSerializer): + class Meta: + model = Spectrogram + fields = '__all__' + + +class SpectrogramViewSet( + mixins.ListModelMixin, + viewsets.GenericViewSet, + LargeImageFileDetailMixin, +): + queryset = Spectrogram.objects.all() + serializer_class = SpectrogramSerializer + + FILE_FIELD_NAME = 'image_file' diff --git a/bats_ai/core/tasks.py b/bats_ai/core/tasks.py index a2ffdc3..8a7780c 100644 --- a/bats_ai/core/tasks.py +++ b/bats_ai/core/tasks.py @@ -1,6 +1,6 @@ from celery import shared_task -from bats_ai.core.models import Image, Recording +from bats_ai.core.models import Image, Recording, colormap @shared_task @@ -13,6 +13,14 @@ def image_compute_checksum(image_id: int): @shared_task def recording_compute_spectrogram(recording_id: int): recording = Recording.objects.get(pk=recording_id) - if not recording.has_spectrogram: - recording.spectrogram # compute by simply referenceing the attribute - assert recording.has_spectrogram + + cmaps = [ + None, # Default (dark) spectrogram + 'light', # Light spectrogram + ] + + for cmap in cmaps: + with colormap(cmap): + if not recording.has_spectrogram: + recording.spectrogram # compute by simply referenceing the attribute + assert recording.has_spectrogram diff --git a/bats_ai/core/views/recording.py b/bats_ai/core/views/recording.py index 2752215..3f2c60b 100644 --- a/bats_ai/core/views/recording.py +++ b/bats_ai/core/views/recording.py @@ -11,7 +11,7 @@ from ninja.files import UploadedFile from ninja.pagination import RouterPaginated -from bats_ai.core.models import Annotations, Recording, Species, TemporalAnnotations +from bats_ai.core.models import Annotations, Recording, Species, TemporalAnnotations, colormap from bats_ai.core.tasks import recording_compute_spectrogram from bats_ai.core.views.species import SpeciesSchema from bats_ai.core.views.temporal_annotations import ( @@ -211,7 +211,8 @@ def get_spectrogram(request: HttpRequest, id: int): except Recording.DoesNotExist: return {'error': 'Recording not found'} - spectrogram = recording.spectrogram + with colormap('light'): + spectrogram = recording.spectrogram spectro_data = { 'base64_spectrogram': spectrogram.base64, @@ -273,22 +274,27 @@ def get_spectrogram_compressed(request: HttpRequest, id: int): except Recording.DoesNotExist: return {'error': 'Recording not found'} - spectrogram = recording.spectrogram - compressed, starts, ends, widths, total_width = spectrogram.compressed + with colormap(): + label, score, confs = recording.spectrogram.predict() + print(label, score, confs) + + with colormap('light'): + spectrogram = recording.spectrogram + _, compressed_base64, metadata = spectrogram.compressed spectro_data = { - 'base64_spectrogram': compressed, + 'base64_spectrogram': compressed_base64, 'spectroInfo': { 'width': spectrogram.width, 'start_time': 0, 'end_time': spectrogram.duration, 'height': spectrogram.height, - 'start_times': starts, - 'end_times': ends, 'low_freq': spectrogram.frequency_min, 'high_freq': spectrogram.frequency_max, - 'compressedWidth': total_width, - 'widths': widths, + 'start_times': metadata['starts'], + 'end_times': metadata['stops'], + 'widths': metadata['widths'], + 'compressedWidth': metadata['length'], }, } diff --git a/bats_ai/settings.py b/bats_ai/settings.py index 11c29dd..7e625da 100644 --- a/bats_ai/settings.py +++ b/bats_ai/settings.py @@ -31,6 +31,7 @@ def mutate_configuration(configuration: ComposedConfiguration) -> None: # Install additional apps configuration.INSTALLED_APPS += [ 'django.contrib.gis', + 'django_large_image', ] configuration.MIDDLEWARE = [ diff --git a/bats_ai/urls.py b/bats_ai/urls.py index 4801c07..5a684a8 100644 --- a/bats_ai/urls.py +++ b/bats_ai/urls.py @@ -3,15 +3,12 @@ from django.urls import include, path from drf_yasg import openapi from drf_yasg.views import get_schema_view -from rest_framework import permissions, routers +from rest_framework import permissions -from bats_ai.core.rest import ImageViewSet +from bats_ai.core.rest import rest from .api import api -router = routers.SimpleRouter() - - # Some more specific Api Requests # OpenAPI generation schema_view = get_schema_view( @@ -19,17 +16,17 @@ public=True, permission_classes=(permissions.AllowAny,), ) -router.register(r'images', ImageViewSet) urlpatterns = [ path('accounts/', include('allauth.urls')), path('oauth/', include('oauth2_provider.urls')), path('admin/', admin.site.urls), path('api/v1/s3-upload/', include('s3_file_field.urls')), - path('api/v1/', include(router.urls)), + path('api/v1/dynamic/', include(rest.urls)), + path('api/v1/', api.urls), path('api/docs/redoc/', schema_view.with_ui('redoc'), name='docs-redoc'), path('api/docs/swagger/', schema_view.with_ui('swagger'), name='docs-swagger'), - path('api/v1/', api.urls), + path('', include('django_large_image.urls')), ] if settings.DEBUG: diff --git a/client/yarn.lock b/client/yarn.lock index fe76a4f..e10d6ce 100644 --- a/client/yarn.lock +++ b/client/yarn.lock @@ -3074,10 +3074,10 @@ es-module-lexer@^1.2.1: resolved "https://registry.npmjs.org/es-module-lexer/-/es-module-lexer-1.4.1.tgz" integrity sha512-cXLGjP0c4T3flZJKQSuziYoq7MlT+rnvfZjfp7h+I7K9BNX54kP9nyWvdbwjQ4u1iWbOL4u96fgeZLToQlZC7w== -esbuild-linux-64@0.14.54: +esbuild-darwin-arm64@0.14.54: version "0.14.54" - resolved "https://registry.npmjs.org/esbuild-linux-64/-/esbuild-linux-64-0.14.54.tgz" - integrity sha512-EgjAgH5HwTbtNsTqQOXWApBaPVdDn7XcK+/PtJwZLT1UmpLoznPd8c5CxqsH2dQK3j05YsB3L17T8vE7cp4cCg== + resolved "https://registry.npmjs.org/esbuild-darwin-arm64/-/esbuild-darwin-arm64-0.14.54.tgz" + integrity sha512-OPafJHD2oUPyvJMrsCvDGkRrVCar5aVyHfWGQzY1dWnzErjrDuSETxwA2HSsyg2jORLY8yBfzc1MIpUkXlctmw== esbuild@^0.14.27: version "0.14.54" @@ -3440,6 +3440,11 @@ fs.realpath@^1.0.0: resolved "https://registry.npmjs.org/fs.realpath/-/fs.realpath-1.0.0.tgz" integrity sha1-FQStJSMVjKpA20onh8sBQRmU6k8= sha512-OO0pH2lK6a0hZnAdau5ItzHPI6pUlvI7jMVnxUQRtw4owF2wk8lOSabtGDCTP4Ggrg2MbGnWO9X8K1t4+fGMDw== +fsevents@^2.3.2, fsevents@~2.3.2: + version "2.3.2" + resolved "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz" + integrity sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA== + function-bind@^1.1.1, function-bind@^1.1.2: version "1.1.2" resolved "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz" diff --git a/dev/django.Dockerfile b/dev/django.Dockerfile index 7fa3378..5992260 100644 --- a/dev/django.Dockerfile +++ b/dev/django.Dockerfile @@ -11,8 +11,10 @@ RUN set -ex \ && apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ gcc \ + g++ \ libc6-dev \ libgdal32 \ + libgdal-dev \ libpq-dev \ libsndfile1-dev \ && apt-get clean \ diff --git a/setup.py b/setup.py index 6589de7..da46ce0 100644 --- a/setup.py +++ b/setup.py @@ -56,7 +56,14 @@ 'librosa', 'matplotlib', 'numpy', + 'onnx', + 'onnxruntime', + # 'onnxruntime-gpu', 'opencv-python-headless', + 'tqdm', + 'django-large-image', + 'large-image[rasterio,pil]>=1.22', + 'rio-cogeo' 'mercantile', ], extras_require={ 'dev': [ From 5cd59288c591c78dae98326add1ada2721bf0e41 Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Mon, 1 Apr 2024 12:17:22 -0400 Subject: [PATCH 2/8] requirements fix --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index da46ce0..8a9139d 100644 --- a/setup.py +++ b/setup.py @@ -63,7 +63,8 @@ 'tqdm', 'django-large-image', 'large-image[rasterio,pil]>=1.22', - 'rio-cogeo' 'mercantile', + 'rio-cogeo', + 'mercantile', ], extras_require={ 'dev': [ From 42a1572dda4bc908d5208e7d8f3be9fc52383fc2 Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Wed, 10 Apr 2024 15:54:56 -0400 Subject: [PATCH 3/8] update spectro --- .../migrations/0011_spectrogram_colormap.py | 18 ++++++++++++++++++ bats_ai/core/models/recording.py | 1 + bats_ai/core/models/spectrogram.py | 4 ++-- bats_ai/core/views/guanometadata.py | 3 ++- bats_ai/core/views/recording.py | 4 ++-- 5 files changed, 25 insertions(+), 5 deletions(-) create mode 100644 bats_ai/core/migrations/0011_spectrogram_colormap.py diff --git a/bats_ai/core/migrations/0011_spectrogram_colormap.py b/bats_ai/core/migrations/0011_spectrogram_colormap.py new file mode 100644 index 0000000..9052b5d --- /dev/null +++ b/bats_ai/core/migrations/0011_spectrogram_colormap.py @@ -0,0 +1,18 @@ +# Generated by Django 4.1.13 on 2024-04-09 13:49 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0010_recording_computed_species_recording_detector_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='spectrogram', + name='colormap', + field=models.CharField(max_length=20, null=True), + ), + ] diff --git a/bats_ai/core/models/recording.py b/bats_ai/core/models/recording.py index 4b63650..15e3288 100644 --- a/bats_ai/core/models/recording.py +++ b/bats_ai/core/models/recording.py @@ -3,6 +3,7 @@ from django.contrib.auth.models import User from django.contrib.gis.db import models from django_extensions.db.models import TimeStampedModel + from .species import Species logger = logging.getLogger(__name__) diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index c6892e7..4640923 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -150,9 +150,9 @@ def generate(cls, recording, colormap=None, dpi=600): img = img.resize((w_, h_), resample=Image.Resampling.LANCZOS) w, h = img.size - img.save(f'temp.{colormap}.png') + img.save(f'/tmp/temp.{colormap}.jpg') - img_filtered = cv2.imread(f'temp.{colormap}.jpg', cv2.IMREAD_GRAYSCALE) + img_filtered = cv2.imread(f'/tmp/temp.{colormap}.jpg', cv2.IMREAD_GRAYSCALE) img_filtered = img_filtered.astype(np.float32) assert img_filtered.min() == 0 assert img_filtered.max() == 255 diff --git a/bats_ai/core/views/guanometadata.py b/bats_ai/core/views/guanometadata.py index bf7e4ed..0dc6d6e 100644 --- a/bats_ai/core/views/guanometadata.py +++ b/bats_ai/core/views/guanometadata.py @@ -44,7 +44,8 @@ def default_data( 'nabat_longitude': (gfile.get('NABat|Longitude', None)), 'nabat_site_name': gfile.get('NABat|Site Name', None), } - + if nabat_fields['nabat_longitude'] and nabat_fields['nabat_longitude'] > 0: # individuals don't put the - in the longitude + nabat_fields['nabat_longitude'] = str((float(nabat_fields['nabat_longitude']) * -1)) # Extract additional fields with conditionals additional_fields = { 'nabat_activation_start_time': parse_datetime( diff --git a/bats_ai/core/views/recording.py b/bats_ai/core/views/recording.py index 3e48239..17be8df 100644 --- a/bats_ai/core/views/recording.py +++ b/bats_ai/core/views/recording.py @@ -268,7 +268,7 @@ def get_spectrogram(request: HttpRequest, id: int): except Recording.DoesNotExist: return {'error': 'Recording not found'} - with colormap('light'): + with colormap(None): spectrogram = recording.spectrogram spectro_data = { @@ -335,7 +335,7 @@ def get_spectrogram_compressed(request: HttpRequest, id: int): label, score, confs = recording.spectrogram.predict() print(label, score, confs) - with colormap('light'): + with colormap(None): spectrogram = recording.spectrogram _, compressed_base64, metadata = spectrogram.compressed From 4c98154dfda3633d30d847d9b87c4d39d161d0cc Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Fri, 12 Apr 2024 16:12:52 -0400 Subject: [PATCH 4/8] compressed spectrogram model --- bats_ai/core/admin/__init__.py | 2 + bats_ai/core/admin/compressed_spectrogram.py | 29 +++ bats_ai/core/admin/spectrogram.py | 14 ++ .../core/migrations/0009_annotations_type.py | 17 -- ...pe_recording_computed_species_and_more.py} | 23 ++- .../migrations/0009_spectrogram_colormap.py | 17 -- .../migrations/0010_compressedspectrogram.py | 36 ++++ .../migrations/0011_spectrogram_colormap.py | 18 -- bats_ai/core/models/__init__.py | 2 + bats_ai/core/models/compressed_spectrogram.py | 116 +++++++++++ bats_ai/core/models/recording.py | 12 +- bats_ai/core/models/spectrogram.py | 153 +------------- bats_ai/core/tasks.py | 190 +++++++++++++++++- 13 files changed, 414 insertions(+), 215 deletions(-) create mode 100644 bats_ai/core/admin/compressed_spectrogram.py delete mode 100644 bats_ai/core/migrations/0009_annotations_type.py rename bats_ai/core/migrations/{0010_recording_computed_species_recording_detector_and_more.py => 0009_annotations_type_recording_computed_species_and_more.py} (64%) delete mode 100644 bats_ai/core/migrations/0009_spectrogram_colormap.py create mode 100644 bats_ai/core/migrations/0010_compressedspectrogram.py delete mode 100644 bats_ai/core/migrations/0011_spectrogram_colormap.py create mode 100644 bats_ai/core/models/compressed_spectrogram.py diff --git a/bats_ai/core/admin/__init__.py b/bats_ai/core/admin/__init__.py index 48b05ad..728cc62 100644 --- a/bats_ai/core/admin/__init__.py +++ b/bats_ai/core/admin/__init__.py @@ -5,6 +5,7 @@ from .species import SpeciesAdmin from .spectrogram import SpectrogramAdmin from .temporal_annotations import TemporalAnnotationsAdmin +from .compressed_spectrogram import CompressedSpectrogramAdmin __all__ = [ 'AnnotationsAdmin', @@ -14,4 +15,5 @@ 'TemporalAnnotationsAdmin', 'SpeciesAdmin', 'GRTSCellsAdmin', + 'CompressedSpectrogramAdmin', ] diff --git a/bats_ai/core/admin/compressed_spectrogram.py b/bats_ai/core/admin/compressed_spectrogram.py new file mode 100644 index 0000000..9dd6cd8 --- /dev/null +++ b/bats_ai/core/admin/compressed_spectrogram.py @@ -0,0 +1,29 @@ +from django.contrib import admin + +from bats_ai.core.models import CompressedSpectrogram + +@admin.register(CompressedSpectrogram) +class CompressedSpectrogramAdmin(admin.ModelAdmin): + list_display = [ + 'pk', + 'recording', + 'spectrogram', + 'image_file', + 'length', + 'widths', + 'starts', + 'stops', + ] + list_display_links = ['pk', 'recording', 'spectrogram'] + list_select_related = True + autocomplete_fields = ['recording'] + readonly_fields = [ + 'recording', + 'spectrogram', + 'image_file', + 'created', + 'modified', + 'widths', + 'starts', + 'stops', + ] diff --git a/bats_ai/core/admin/spectrogram.py b/bats_ai/core/admin/spectrogram.py index 6f2f601..23bfc15 100644 --- a/bats_ai/core/admin/spectrogram.py +++ b/bats_ai/core/admin/spectrogram.py @@ -1,6 +1,8 @@ from django.contrib import admin from bats_ai.core.models import Spectrogram +from django.db.models import QuerySet +from django.http import HttpRequest @admin.register(Spectrogram) @@ -31,3 +33,15 @@ class SpectrogramAdmin(admin.ModelAdmin): 'frequency_min', 'frequency_max', ] + + actions = ['computed_compressed_spectrogram'] + + + @admin.action(description='Compute Compressed Spectrograms') + def computed_compressed_spectrogram(self, request: HttpRequest, queryset: QuerySet): + counter = 0 + for recording in queryset: + if not recording.has_spectrogram: + recording_compute_spectrogram.delay(recording.pk) + counter += 1 + self.message_user(request, f'{counter} recordings queued', messages.SUCCESS) diff --git a/bats_ai/core/migrations/0009_annotations_type.py b/bats_ai/core/migrations/0009_annotations_type.py deleted file mode 100644 index 6091cd7..0000000 --- a/bats_ai/core/migrations/0009_annotations_type.py +++ /dev/null @@ -1,17 +0,0 @@ -# Generated by Django 4.1.13 on 2024-03-22 17:23 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ('core', '0008_grtscells_recording_recorded_time'), - ] - - operations = [ - migrations.AddField( - model_name='annotations', - name='type', - field=models.TextField(blank=True, null=True), - ), - ] diff --git a/bats_ai/core/migrations/0010_recording_computed_species_recording_detector_and_more.py b/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py similarity index 64% rename from bats_ai/core/migrations/0010_recording_computed_species_recording_detector_and_more.py rename to bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py index 1819063..004d224 100644 --- a/bats_ai/core/migrations/0010_recording_computed_species_recording_detector_and_more.py +++ b/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py @@ -1,20 +1,24 @@ -# Generated by Django 4.1.13 on 2024-04-03 13:07 +# Generated by Django 4.1.13 on 2024-04-11 13:06 from django.db import migrations, models class Migration(migrations.Migration): + dependencies = [ - ('core', '0009_annotations_type'), + ('core', '0008_grtscells_recording_recorded_time'), ] operations = [ + migrations.AddField( + model_name='annotations', + name='type', + field=models.TextField(blank=True, null=True), + ), migrations.AddField( model_name='recording', name='computed_species', - field=models.ManyToManyField( - related_name='recording_computed_species', to='core.species' - ), + field=models.ManyToManyField(related_name='recording_computed_species', to='core.species'), ), migrations.AddField( model_name='recording', @@ -24,9 +28,7 @@ class Migration(migrations.Migration): migrations.AddField( model_name='recording', name='official_species', - field=models.ManyToManyField( - related_name='recording_official_species', to='core.species' - ), + field=models.ManyToManyField(related_name='recording_official_species', to='core.species'), ), migrations.AddField( model_name='recording', @@ -48,4 +50,9 @@ class Migration(migrations.Migration): name='unusual_occurrences', field=models.TextField(blank=True, null=True), ), + migrations.AddField( + model_name='spectrogram', + name='colormap', + field=models.CharField(max_length=20, null=True), + ), ] diff --git a/bats_ai/core/migrations/0009_spectrogram_colormap.py b/bats_ai/core/migrations/0009_spectrogram_colormap.py deleted file mode 100644 index a89259e..0000000 --- a/bats_ai/core/migrations/0009_spectrogram_colormap.py +++ /dev/null @@ -1,17 +0,0 @@ -# Generated by Django 4.1.13 on 2024-03-12 21:53 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ('core', '0008_grtscells_recording_recorded_time'), - ] - - operations = [ - migrations.AddField( - model_name='spectrogram', - name='colormap', - field=models.CharField(max_length=20, null=True), - ), - ] diff --git a/bats_ai/core/migrations/0010_compressedspectrogram.py b/bats_ai/core/migrations/0010_compressedspectrogram.py new file mode 100644 index 0000000..f7a1fb3 --- /dev/null +++ b/bats_ai/core/migrations/0010_compressedspectrogram.py @@ -0,0 +1,36 @@ +# Generated by Django 4.1.13 on 2024-04-12 18:56 + +import django.contrib.postgres.fields +from django.db import migrations, models +import django.db.models.deletion +import django_extensions.db.fields + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0009_annotations_type_recording_computed_species_and_more'), + ] + + operations = [ + migrations.CreateModel( + name='CompressedSpectrogram', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('created', django_extensions.db.fields.CreationDateTimeField(auto_now_add=True, verbose_name='created')), + ('modified', django_extensions.db.fields.ModificationDateTimeField(auto_now=True, verbose_name='modified')), + ('image_file', models.FileField(upload_to='')), + ('length', models.IntegerField()), + ('starts', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), + ('stops', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), + ('widths', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), + ('cache_invalidated', models.BooleanField(default=True)), + ('recording', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.recording')), + ('spectrogram', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.spectrogram')), + ], + options={ + 'get_latest_by': 'modified', + 'abstract': False, + }, + ), + ] diff --git a/bats_ai/core/migrations/0011_spectrogram_colormap.py b/bats_ai/core/migrations/0011_spectrogram_colormap.py deleted file mode 100644 index 9052b5d..0000000 --- a/bats_ai/core/migrations/0011_spectrogram_colormap.py +++ /dev/null @@ -1,18 +0,0 @@ -# Generated by Django 4.1.13 on 2024-04-09 13:49 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - - dependencies = [ - ('core', '0010_recording_computed_species_recording_detector_and_more'), - ] - - operations = [ - migrations.AddField( - model_name='spectrogram', - name='colormap', - field=models.CharField(max_length=20, null=True), - ), - ] diff --git a/bats_ai/core/models/__init__.py b/bats_ai/core/models/__init__.py index 5979660..ac6f233 100644 --- a/bats_ai/core/models/__init__.py +++ b/bats_ai/core/models/__init__.py @@ -6,6 +6,7 @@ from .species import Species from .spectrogram import Spectrogram from .temporal_annotations import TemporalAnnotations +from .compressed_spectrogram import CompressedSpectrogram __all__ = [ 'Annotations', @@ -17,4 +18,5 @@ 'TemporalAnnotations', 'GRTSCells', 'colormap', + 'CompressedSpectrogram', ] diff --git a/bats_ai/core/models/compressed_spectrogram.py b/bats_ai/core/models/compressed_spectrogram.py new file mode 100644 index 0000000..aa3dadd --- /dev/null +++ b/bats_ai/core/models/compressed_spectrogram.py @@ -0,0 +1,116 @@ +import base64 +import io +import math + +from PIL import Image +import cv2 +from django.core.files import File +from django.db import models +from django.db.models.fields.files import FieldFile +from django.contrib.postgres.fields import ArrayField +from django_extensions.db.models import TimeStampedModel +import librosa +import matplotlib.pyplot as plt +import numpy as np +import scipy +import tqdm + +from bats_ai.core.models import Annotations, Recording, Spectrogram +from django.dispatch import receiver + +FREQ_MIN = 5e3 +FREQ_MAX = 120e3 +FREQ_PAD = 2e3 + + +# TimeStampedModel also provides "created" and "modified" fields +class CompressedSpectrogram(TimeStampedModel, models.Model): + recording = models.ForeignKey(Recording, on_delete=models.CASCADE) + spectrogram = models.ForeignKey(Spectrogram, on_delete=models.CASCADE) + image_file = models.FileField() + length = models.IntegerField() # pixels + starts = ArrayField(ArrayField(models.IntegerField())) # pixels + stops = ArrayField(ArrayField(models.IntegerField())) # milliseconds + widths = ArrayField(ArrayField(models.IntegerField())) # hz + cache_invalidated = models.BooleanField(default=True) + + + def predict(self): + import json + import os + + import onnx + import onnxruntime as ort + import tqdm + + img = Image.open(self.image_file) + + relative = ('..',) * 4 + asset_path = os.path.abspath(os.path.join(__file__, *relative, 'assets')) + + onnx_filename = os.path.join(asset_path, 'model.mobilenet.onnx') + assert os.path.exists(onnx_filename) + + session = ort.InferenceSession( + onnx_filename, + providers=[ + ( + 'CUDAExecutionProvider', + { + 'cudnn_conv_use_max_workspace': '1', + 'device_id': 0, + 'cudnn_conv_algo_search': 'HEURISTIC', + }, + ), + 'CPUExecutionProvider', + ], + ) + + img = np.array(img) + + h, w, c = img.shape + ratio_y = 224 / h + ratio_x = ratio_y * 0.5 + raw = cv2.resize(img, None, fx=ratio_x, fy=ratio_y, interpolation=cv2.INTER_LANCZOS4) + + h, w, c = raw.shape + if w <= h: + canvas = np.zeros((h, h + 1, 3), dtype=raw.dtype) + canvas[:, :w, :] = raw + raw = canvas + h, w, c = raw.shape + + inputs_ = [] + for index in range(0, w - h, 100): + inputs_.append(raw[:, index : index + h, :]) + inputs_.append(raw[:, -h:, :]) + inputs_ = np.array(inputs_) + + chunksize = 1 + chunks = np.array_split(inputs_, np.arange(chunksize, len(inputs_), chunksize)) + outputs = [] + for chunk in tqdm.tqdm(chunks, desc='Inference'): + outputs_ = session.run( + None, + {'input': chunk}, + ) + outputs.append(outputs_[0]) + outputs = np.vstack(outputs) + outputs = outputs.mean(axis=0) + + model = onnx.load(onnx_filename) + mapping = json.loads(model.metadata_props[0].value) + labels = [mapping['forward'][str(index)] for index in range(len(mapping['forward']))] + + prediction = np.argmax(outputs) + label = labels[prediction] + score = outputs[prediction] + + confs = dict(zip(labels, outputs)) + + return label, score, confs + +@receiver(models.signals.pre_delete, sender=Spectrogram) +def delete_content(sender, instance, **kwargs): + if instance.image_file: + instance.image_file.delete(save=False) diff --git a/bats_ai/core/models/recording.py b/bats_ai/core/models/recording.py index 15e3288..78b3e59 100644 --- a/bats_ai/core/models/recording.py +++ b/bats_ai/core/models/recording.py @@ -3,6 +3,7 @@ from django.contrib.auth.models import User from django.contrib.gis.db import models from django_extensions.db.models import TimeStampedModel +from django.dispatch import receiver from .species import Species @@ -86,13 +87,12 @@ def spectrogram(self): spectrograms = self.spectrograms - if len(spectrograms) == 0: - Spectrogram.generate(self, colormap=COLORMAP) - - spectrograms = self.spectrograms - assert len(spectrograms) == 1 - assert len(spectrograms) >= 1 spectrogram = spectrograms[0] # most recently created return spectrogram + +@receiver(models.signals.pre_delete, sender=Recording) +def delete_content(sender, instance, **kwargs): + if instance.audio_file: + instance.audio_file.delete(save=False) diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index 4640923..9ff800f 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -15,6 +15,7 @@ import tqdm from bats_ai.core.models import Annotations, Recording +from django.dispatch import receiver FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -155,7 +156,6 @@ def generate(cls, recording, colormap=None, dpi=600): img_filtered = cv2.imread(f'/tmp/temp.{colormap}.jpg', cv2.IMREAD_GRAYSCALE) img_filtered = img_filtered.astype(np.float32) assert img_filtered.min() == 0 - assert img_filtered.max() == 255 img_filtered = img_filtered.astype(np.float32) img_filtered /= 0.9 img_filtered[img_filtered < 0] = 0 @@ -188,7 +188,7 @@ def generate(cls, recording, colormap=None, dpi=600): img.save(buf, format='JPEG', quality=80) buf.seek(0) - name = 'spectrogram.jpg' + name = f'{recording.pk}_{colormap}_spectrogram.jpg' image_file = File(buf, name=name) spectrogram = cls( @@ -202,148 +202,7 @@ def generate(cls, recording, colormap=None, dpi=600): colormap=colormap, ) spectrogram.save() - - @property - def compressed(self): - img = self.image_np - - annotations = Annotations.objects.filter(recording=self.recording) - - threshold = 0.5 - while True: - canvas = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) - canvas = canvas.astype(np.float32) - - is_light = np.median(canvas) > 128.0 - if is_light: - canvas = 255.0 - canvas - - amplitude = canvas.max(axis=0) - amplitude -= amplitude.min() - amplitude /= amplitude.max() - amplitude[amplitude < threshold] = 0.0 - amplitude[amplitude > 0] = 1.0 - amplitude = amplitude.reshape(1, -1) - - canvas -= canvas.min() - canvas /= canvas.max() - canvas *= 255.0 - canvas *= amplitude - canvas = np.around(canvas).astype(np.uint8) - - width = canvas.shape[1] - for annotation in annotations: - start = annotation.start_time / self.duration - stop = annotation.end_time / self.duration - - start = int(np.around(start * width)) - stop = int(np.around(stop * width)) - canvas[:, start : stop + 1] = 255.0 - - mask = canvas.max(axis=0) - mask = scipy.signal.medfilt(mask, 3) - mask[0] = 0 - mask[-1] = 0 - - starts = [] - stops = [] - for index in range(1, len(mask) - 1): - value_pre = mask[index - 1] - value = mask[index] - value_post = mask[index + 1] - if value != 0: - if value_pre == 0: - starts.append(index) - if value_post == 0: - stops.append(index) - assert len(starts) == len(stops) - - starts = [val - 40 for val in starts] # 10 ms buffer - stops = [val + 40 for val in stops] # 10 ms buffer - ranges = list(zip(starts, stops)) - - while True: - found = False - merged = [] - index = 0 - while index < len(ranges) - 1: - start1, stop1 = ranges[index] - start2, stop2 = ranges[index + 1] - - start1 = min(max(start1, 0), len(mask)) - start2 = min(max(start2, 0), len(mask)) - stop1 = min(max(stop1, 0), len(mask)) - stop2 = min(max(stop2, 0), len(mask)) - - if stop1 >= start2: - found = True - merged.append((start1, stop2)) - index += 2 - else: - merged.append((start1, stop1)) - index += 1 - if index == len(ranges) - 1: - merged.append((start2, stop2)) - ranges = merged - if not found: - for index in range(1, len(ranges)): - start1, stop1 = ranges[index - 1] - start2, stop2 = ranges[index] - assert start1 < stop1 - assert start2 < stop2 - assert start1 < start2 - assert stop1 < stop2 - assert stop1 < start2 - break - - segments = [] - starts_ = [] - stops_ = [] - domain = img.shape[1] - widths = [] - total_width = 0 - for start, stop in ranges: - segment = img[:, start:stop] - segments.append(segment) - - starts_.append(int(round(self.duration * (start / domain)))) - stops_.append(int(round(self.duration * (stop / domain)))) - widths.append(stop - start) - total_width += stop - start - - # buffer = np.zeros((len(img), 20, 3), dtype=img.dtype) - # segments.append(buffer) - # segments = segments[:-1] - - if len(segments) > 0: - break - - threshold -= 0.05 - if threshold < 0: - segments = None - break - - if segments is None: - canvas = img.copy() - else: - canvas = np.hstack(segments) - - canvas = Image.fromarray(canvas, 'RGB') - - canvas.save('temp.compressed.jpg') - - buf = io.BytesIO() - canvas.save(buf, format='JPEG', quality=80) - buf.seek(0) - img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8') - - metadata = { - 'starts': starts_, - 'stops': stops_, - 'widths': widths, - 'length': total_width, - } - return canvas, img_base64, metadata + return spectrogram.pk @property def image_np(self): @@ -439,3 +298,9 @@ def predict(self): confs = dict(zip(labels, outputs)) return label, score, confs + +@receiver(models.signals.pre_delete, sender=Spectrogram) +def delete_content(sender, instance, **kwargs): + if instance.image_file: + instance.image_file.delete(save=False) + diff --git a/bats_ai/core/tasks.py b/bats_ai/core/tasks.py index 8a7780c..215e92d 100644 --- a/bats_ai/core/tasks.py +++ b/bats_ai/core/tasks.py @@ -1,6 +1,154 @@ from celery import shared_task -from bats_ai.core.models import Image, Recording, colormap +from bats_ai.core.models import Image, Recording, colormap, Spectrogram, CompressedSpectrogram, Annotations + +import io +import math +from django.core.files import File +from PIL import Image +import cv2 +import librosa +import matplotlib.pyplot as plt +import numpy as np +import scipy +import tqdm + +def generate_compressed(recording: Recording, spectrogram: Spectrogram): + print(spectrogram) + img = spectrogram.image_np + + annotations = Annotations.objects.filter(recording=recording) + + threshold = 0.5 + while True: + canvas = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) + canvas = canvas.astype(np.float32) + + is_light = np.median(canvas) > 128.0 + if is_light: + canvas = 255.0 - canvas + + amplitude = canvas.max(axis=0) + amplitude -= amplitude.min() + amplitude /= amplitude.max() + amplitude[amplitude < threshold] = 0.0 + amplitude[amplitude > 0] = 1.0 + amplitude = amplitude.reshape(1, -1) + + canvas -= canvas.min() + canvas /= canvas.max() + canvas *= 255.0 + canvas *= amplitude + canvas = np.around(canvas).astype(np.uint8) + + width = canvas.shape[1] + for annotation in annotations: + start = annotation.start_time / spectrogram.duration + stop = annotation.end_time / spectrogram.duration + + start = int(np.around(start * width)) + stop = int(np.around(stop * width)) + canvas[:, start : stop + 1] = 255.0 + + mask = canvas.max(axis=0) + mask = scipy.signal.medfilt(mask, 3) + mask[0] = 0 + mask[-1] = 0 + + starts = [] + stops = [] + for index in range(1, len(mask) - 1): + value_pre = mask[index - 1] + value = mask[index] + value_post = mask[index + 1] + if value != 0: + if value_pre == 0: + starts.append(index) + if value_post == 0: + stops.append(index) + assert len(starts) == len(stops) + + starts = [val - 40 for val in starts] # 10 ms buffer + stops = [val + 40 for val in stops] # 10 ms buffer + ranges = list(zip(starts, stops)) + + while True: + found = False + merged = [] + index = 0 + while index < len(ranges) - 1: + start1, stop1 = ranges[index] + start2, stop2 = ranges[index + 1] + + start1 = min(max(start1, 0), len(mask)) + start2 = min(max(start2, 0), len(mask)) + stop1 = min(max(stop1, 0), len(mask)) + stop2 = min(max(stop2, 0), len(mask)) + + if stop1 >= start2: + found = True + merged.append((start1, stop2)) + index += 2 + else: + merged.append((start1, stop1)) + index += 1 + if index == len(ranges) - 1: + merged.append((start2, stop2)) + ranges = merged + if not found: + for index in range(1, len(ranges)): + start1, stop1 = ranges[index - 1] + start2, stop2 = ranges[index] + assert start1 < stop1 + assert start2 < stop2 + assert start1 < start2 + assert stop1 < stop2 + assert stop1 < start2 + break + + segments = [] + starts_ = [] + stops_ = [] + domain = img.shape[1] + widths = [] + total_width = 0 + for start, stop in ranges: + segment = img[:, start:stop] + segments.append(segment) + + starts_.append(int(round(spectrogram.duration * (start / domain)))) + stops_.append(int(round(spectrogram.duration * (stop / domain)))) + widths.append(stop - start) + total_width += stop - start + + # buffer = np.zeros((len(img), 20, 3), dtype=img.dtype) + # segments.append(buffer) + # segments = segments[:-1] + + if len(segments) > 0: + break + + threshold -= 0.05 + if threshold < 0: + segments = None + break + + if segments is None: + canvas = img.copy() + else: + canvas = np.hstack(segments) + + canvas = Image.fromarray(canvas, 'RGB') + + canvas.save('temp.compressed.jpg') + + buf = io.BytesIO() + canvas.save(buf, format='JPEG', quality=80) + buf.seek(0) + name = f'{spectrogram.pk}_spectrogram_compressed.jpg' + image_file = File(buf, name=name) + + return total_width, image_file, widths, starts, stops @shared_task @@ -18,9 +166,41 @@ def recording_compute_spectrogram(recording_id: int): None, # Default (dark) spectrogram 'light', # Light spectrogram ] - + spectrogram_id = None for cmap in cmaps: with colormap(cmap): - if not recording.has_spectrogram: - recording.spectrogram # compute by simply referenceing the attribute - assert recording.has_spectrogram + spectrogram_id_temp = Spectrogram.generate(recording, cmap) + if cmap is None: + spectrogram_id = spectrogram_id_temp + if spectrogram_id is not None: + generate_compress_spectrogram.delay(recording_id, spectrogram_id) + + +@shared_task +def generate_compress_spectrogram(recording_id: int, spectrogram_id: int): + recording = Recording.objects.get(pk=recording_id) + spectrogram = Spectrogram.objects.get(pk=spectrogram_id) + length, image_file, widths, starts, stops = generate_compressed(recording, spectrogram) + found = CompressedSpectrogram.objects.filter(recording=recording, spectrogram=spectrogram) + if found.exists(): + existing = found.first() + existing.length = length + existing.image_file = image_file + existing.widths = widths + existing.starts = starts + existing.stops = stops + existing.cache_invalidated = False + existing.save() + else: + CompressedSpectrogram.objects.create( + recording=recording, + spectrogram=spectrogram, + image_file=image_file, + length=length, + widths=widths, + starts=starts, + stops=stops, + cache_invalidated=False + ) + + From 4e1035fd94a73c7481fa4e811e5f28b95ab62e9e Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Fri, 12 Apr 2024 16:16:16 -0400 Subject: [PATCH 5/8] linting --- bats_ai/core/admin/__init__.py | 2 +- bats_ai/core/admin/compressed_spectrogram.py | 1 + bats_ai/core/admin/spectrogram.py | 14 - ...ype_recording_computed_species_and_more.py | 9 +- .../migrations/0010_compressedspectrogram.py | 66 ++++- bats_ai/core/models/__init__.py | 2 +- bats_ai/core/models/compressed_spectrogram.py | 18 +- bats_ai/core/models/recording.py | 5 +- bats_ai/core/models/spectrogram.py | 7 +- bats_ai/core/tasks.py | 273 +++++++++--------- 10 files changed, 210 insertions(+), 187 deletions(-) diff --git a/bats_ai/core/admin/__init__.py b/bats_ai/core/admin/__init__.py index 728cc62..8f782b8 100644 --- a/bats_ai/core/admin/__init__.py +++ b/bats_ai/core/admin/__init__.py @@ -1,11 +1,11 @@ from .annotations import AnnotationsAdmin +from .compressed_spectrogram import CompressedSpectrogramAdmin from .grts_cells import GRTSCellsAdmin from .image import ImageAdmin from .recording import RecordingAdmin from .species import SpeciesAdmin from .spectrogram import SpectrogramAdmin from .temporal_annotations import TemporalAnnotationsAdmin -from .compressed_spectrogram import CompressedSpectrogramAdmin __all__ = [ 'AnnotationsAdmin', diff --git a/bats_ai/core/admin/compressed_spectrogram.py b/bats_ai/core/admin/compressed_spectrogram.py index 9dd6cd8..a8ce90c 100644 --- a/bats_ai/core/admin/compressed_spectrogram.py +++ b/bats_ai/core/admin/compressed_spectrogram.py @@ -2,6 +2,7 @@ from bats_ai.core.models import CompressedSpectrogram + @admin.register(CompressedSpectrogram) class CompressedSpectrogramAdmin(admin.ModelAdmin): list_display = [ diff --git a/bats_ai/core/admin/spectrogram.py b/bats_ai/core/admin/spectrogram.py index 23bfc15..6f2f601 100644 --- a/bats_ai/core/admin/spectrogram.py +++ b/bats_ai/core/admin/spectrogram.py @@ -1,8 +1,6 @@ from django.contrib import admin from bats_ai.core.models import Spectrogram -from django.db.models import QuerySet -from django.http import HttpRequest @admin.register(Spectrogram) @@ -33,15 +31,3 @@ class SpectrogramAdmin(admin.ModelAdmin): 'frequency_min', 'frequency_max', ] - - actions = ['computed_compressed_spectrogram'] - - - @admin.action(description='Compute Compressed Spectrograms') - def computed_compressed_spectrogram(self, request: HttpRequest, queryset: QuerySet): - counter = 0 - for recording in queryset: - if not recording.has_spectrogram: - recording_compute_spectrogram.delay(recording.pk) - counter += 1 - self.message_user(request, f'{counter} recordings queued', messages.SUCCESS) diff --git a/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py b/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py index 004d224..ce429a8 100644 --- a/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py +++ b/bats_ai/core/migrations/0009_annotations_type_recording_computed_species_and_more.py @@ -4,7 +4,6 @@ class Migration(migrations.Migration): - dependencies = [ ('core', '0008_grtscells_recording_recorded_time'), ] @@ -18,7 +17,9 @@ class Migration(migrations.Migration): migrations.AddField( model_name='recording', name='computed_species', - field=models.ManyToManyField(related_name='recording_computed_species', to='core.species'), + field=models.ManyToManyField( + related_name='recording_computed_species', to='core.species' + ), ), migrations.AddField( model_name='recording', @@ -28,7 +29,9 @@ class Migration(migrations.Migration): migrations.AddField( model_name='recording', name='official_species', - field=models.ManyToManyField(related_name='recording_official_species', to='core.species'), + field=models.ManyToManyField( + related_name='recording_official_species', to='core.species' + ), ), migrations.AddField( model_name='recording', diff --git a/bats_ai/core/migrations/0010_compressedspectrogram.py b/bats_ai/core/migrations/0010_compressedspectrogram.py index f7a1fb3..0ce5b18 100644 --- a/bats_ai/core/migrations/0010_compressedspectrogram.py +++ b/bats_ai/core/migrations/0010_compressedspectrogram.py @@ -7,7 +7,6 @@ class Migration(migrations.Migration): - dependencies = [ ('core', '0009_annotations_type_recording_computed_species_and_more'), ] @@ -16,17 +15,66 @@ class Migration(migrations.Migration): migrations.CreateModel( name='CompressedSpectrogram', fields=[ - ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('created', django_extensions.db.fields.CreationDateTimeField(auto_now_add=True, verbose_name='created')), - ('modified', django_extensions.db.fields.ModificationDateTimeField(auto_now=True, verbose_name='modified')), + ( + 'id', + models.BigAutoField( + auto_created=True, primary_key=True, serialize=False, verbose_name='ID' + ), + ), + ( + 'created', + django_extensions.db.fields.CreationDateTimeField( + auto_now_add=True, verbose_name='created' + ), + ), + ( + 'modified', + django_extensions.db.fields.ModificationDateTimeField( + auto_now=True, verbose_name='modified' + ), + ), ('image_file', models.FileField(upload_to='')), ('length', models.IntegerField()), - ('starts', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), - ('stops', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), - ('widths', django.contrib.postgres.fields.ArrayField(base_field=django.contrib.postgres.fields.ArrayField(base_field=models.IntegerField(), size=None), size=None)), + ( + 'starts', + django.contrib.postgres.fields.ArrayField( + base_field=django.contrib.postgres.fields.ArrayField( + base_field=models.IntegerField(), size=None + ), + size=None, + ), + ), + ( + 'stops', + django.contrib.postgres.fields.ArrayField( + base_field=django.contrib.postgres.fields.ArrayField( + base_field=models.IntegerField(), size=None + ), + size=None, + ), + ), + ( + 'widths', + django.contrib.postgres.fields.ArrayField( + base_field=django.contrib.postgres.fields.ArrayField( + base_field=models.IntegerField(), size=None + ), + size=None, + ), + ), ('cache_invalidated', models.BooleanField(default=True)), - ('recording', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.recording')), - ('spectrogram', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.spectrogram')), + ( + 'recording', + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, to='core.recording' + ), + ), + ( + 'spectrogram', + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, to='core.spectrogram' + ), + ), ], options={ 'get_latest_by': 'modified', diff --git a/bats_ai/core/models/__init__.py b/bats_ai/core/models/__init__.py index ac6f233..36e2f0f 100644 --- a/bats_ai/core/models/__init__.py +++ b/bats_ai/core/models/__init__.py @@ -1,4 +1,5 @@ from .annotations import Annotations +from .compressed_spectrogram import CompressedSpectrogram from .grts_cells import GRTSCells from .image import Image from .recording import Recording, colormap @@ -6,7 +7,6 @@ from .species import Species from .spectrogram import Spectrogram from .temporal_annotations import TemporalAnnotations -from .compressed_spectrogram import CompressedSpectrogram __all__ = [ 'Annotations', diff --git a/bats_ai/core/models/compressed_spectrogram.py b/bats_ai/core/models/compressed_spectrogram.py index aa3dadd..04466bf 100644 --- a/bats_ai/core/models/compressed_spectrogram.py +++ b/bats_ai/core/models/compressed_spectrogram.py @@ -1,22 +1,12 @@ -import base64 -import io -import math - from PIL import Image import cv2 -from django.core.files import File -from django.db import models -from django.db.models.fields.files import FieldFile from django.contrib.postgres.fields import ArrayField +from django.db import models +from django.dispatch import receiver from django_extensions.db.models import TimeStampedModel -import librosa -import matplotlib.pyplot as plt import numpy as np -import scipy -import tqdm -from bats_ai.core.models import Annotations, Recording, Spectrogram -from django.dispatch import receiver +from bats_ai.core.models import Recording, Spectrogram FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -33,7 +23,6 @@ class CompressedSpectrogram(TimeStampedModel, models.Model): stops = ArrayField(ArrayField(models.IntegerField())) # milliseconds widths = ArrayField(ArrayField(models.IntegerField())) # hz cache_invalidated = models.BooleanField(default=True) - def predict(self): import json @@ -110,6 +99,7 @@ def predict(self): return label, score, confs + @receiver(models.signals.pre_delete, sender=Spectrogram) def delete_content(sender, instance, **kwargs): if instance.image_file: diff --git a/bats_ai/core/models/recording.py b/bats_ai/core/models/recording.py index 78b3e59..25266e9 100644 --- a/bats_ai/core/models/recording.py +++ b/bats_ai/core/models/recording.py @@ -2,8 +2,8 @@ from django.contrib.auth.models import User from django.contrib.gis.db import models -from django_extensions.db.models import TimeStampedModel from django.dispatch import receiver +from django_extensions.db.models import TimeStampedModel from .species import Species @@ -83,7 +83,7 @@ def spectrograms(self): @property def spectrogram(self): - from bats_ai.core.models import Spectrogram + pass spectrograms = self.spectrograms @@ -92,6 +92,7 @@ def spectrogram(self): return spectrogram + @receiver(models.signals.pre_delete, sender=Recording) def delete_content(sender, instance, **kwargs): if instance.audio_file: diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index 9ff800f..a31ffee 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -7,15 +7,14 @@ from django.core.files import File from django.db import models from django.db.models.fields.files import FieldFile +from django.dispatch import receiver from django_extensions.db.models import TimeStampedModel import librosa import matplotlib.pyplot as plt import numpy as np -import scipy import tqdm -from bats_ai.core.models import Annotations, Recording -from django.dispatch import receiver +from bats_ai.core.models import Recording FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -299,8 +298,8 @@ def predict(self): return label, score, confs + @receiver(models.signals.pre_delete, sender=Spectrogram) def delete_content(sender, instance, **kwargs): if instance.image_file: instance.image_file.delete(save=False) - diff --git a/bats_ai/core/tasks.py b/bats_ai/core/tasks.py index 215e92d..b1abd90 100644 --- a/bats_ai/core/tasks.py +++ b/bats_ai/core/tasks.py @@ -1,154 +1,151 @@ -from celery import shared_task - -from bats_ai.core.models import Image, Recording, colormap, Spectrogram, CompressedSpectrogram, Annotations - import io -import math -from django.core.files import File + from PIL import Image +from celery import shared_task import cv2 -import librosa -import matplotlib.pyplot as plt +from django.core.files import File import numpy as np import scipy -import tqdm -def generate_compressed(recording: Recording, spectrogram: Spectrogram): - print(spectrogram) - img = spectrogram.image_np +from bats_ai.core.models import Annotations, CompressedSpectrogram, Recording, Spectrogram, colormap - annotations = Annotations.objects.filter(recording=recording) - threshold = 0.5 - while True: - canvas = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) - canvas = canvas.astype(np.float32) - - is_light = np.median(canvas) > 128.0 - if is_light: - canvas = 255.0 - canvas - - amplitude = canvas.max(axis=0) - amplitude -= amplitude.min() - amplitude /= amplitude.max() - amplitude[amplitude < threshold] = 0.0 - amplitude[amplitude > 0] = 1.0 - amplitude = amplitude.reshape(1, -1) - - canvas -= canvas.min() - canvas /= canvas.max() - canvas *= 255.0 - canvas *= amplitude - canvas = np.around(canvas).astype(np.uint8) - - width = canvas.shape[1] - for annotation in annotations: - start = annotation.start_time / spectrogram.duration - stop = annotation.end_time / spectrogram.duration - - start = int(np.around(start * width)) - stop = int(np.around(stop * width)) - canvas[:, start : stop + 1] = 255.0 - - mask = canvas.max(axis=0) - mask = scipy.signal.medfilt(mask, 3) - mask[0] = 0 - mask[-1] = 0 - - starts = [] - stops = [] - for index in range(1, len(mask) - 1): - value_pre = mask[index - 1] - value = mask[index] - value_post = mask[index + 1] - if value != 0: - if value_pre == 0: - starts.append(index) - if value_post == 0: - stops.append(index) - assert len(starts) == len(stops) - - starts = [val - 40 for val in starts] # 10 ms buffer - stops = [val + 40 for val in stops] # 10 ms buffer - ranges = list(zip(starts, stops)) - - while True: - found = False - merged = [] - index = 0 - while index < len(ranges) - 1: - start1, stop1 = ranges[index] - start2, stop2 = ranges[index + 1] - - start1 = min(max(start1, 0), len(mask)) - start2 = min(max(start2, 0), len(mask)) - stop1 = min(max(stop1, 0), len(mask)) - stop2 = min(max(stop2, 0), len(mask)) - - if stop1 >= start2: - found = True - merged.append((start1, stop2)) - index += 2 - else: - merged.append((start1, stop1)) - index += 1 - if index == len(ranges) - 1: - merged.append((start2, stop2)) - ranges = merged - if not found: - for index in range(1, len(ranges)): - start1, stop1 = ranges[index - 1] - start2, stop2 = ranges[index] - assert start1 < stop1 - assert start2 < stop2 - assert start1 < start2 - assert stop1 < stop2 - assert stop1 < start2 - break - - segments = [] - starts_ = [] - stops_ = [] - domain = img.shape[1] - widths = [] - total_width = 0 - for start, stop in ranges: - segment = img[:, start:stop] - segments.append(segment) - - starts_.append(int(round(spectrogram.duration * (start / domain)))) - stops_.append(int(round(spectrogram.duration * (stop / domain)))) - widths.append(stop - start) - total_width += stop - start - - # buffer = np.zeros((len(img), 20, 3), dtype=img.dtype) - # segments.append(buffer) - # segments = segments[:-1] - - if len(segments) > 0: - break +def generate_compressed(recording: Recording, spectrogram: Spectrogram): + print(spectrogram) + img = spectrogram.image_np + + annotations = Annotations.objects.filter(recording=recording) + + threshold = 0.5 + while True: + canvas = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) + canvas = canvas.astype(np.float32) + + is_light = np.median(canvas) > 128.0 + if is_light: + canvas = 255.0 - canvas + + amplitude = canvas.max(axis=0) + amplitude -= amplitude.min() + amplitude /= amplitude.max() + amplitude[amplitude < threshold] = 0.0 + amplitude[amplitude > 0] = 1.0 + amplitude = amplitude.reshape(1, -1) + + canvas -= canvas.min() + canvas /= canvas.max() + canvas *= 255.0 + canvas *= amplitude + canvas = np.around(canvas).astype(np.uint8) + + width = canvas.shape[1] + for annotation in annotations: + start = annotation.start_time / spectrogram.duration + stop = annotation.end_time / spectrogram.duration + + start = int(np.around(start * width)) + stop = int(np.around(stop * width)) + canvas[:, start : stop + 1] = 255.0 + + mask = canvas.max(axis=0) + mask = scipy.signal.medfilt(mask, 3) + mask[0] = 0 + mask[-1] = 0 + + starts = [] + stops = [] + for index in range(1, len(mask) - 1): + value_pre = mask[index - 1] + value = mask[index] + value_post = mask[index + 1] + if value != 0: + if value_pre == 0: + starts.append(index) + if value_post == 0: + stops.append(index) + assert len(starts) == len(stops) + + starts = [val - 40 for val in starts] # 10 ms buffer + stops = [val + 40 for val in stops] # 10 ms buffer + ranges = list(zip(starts, stops)) - threshold -= 0.05 - if threshold < 0: - segments = None + while True: + found = False + merged = [] + index = 0 + while index < len(ranges) - 1: + start1, stop1 = ranges[index] + start2, stop2 = ranges[index + 1] + + start1 = min(max(start1, 0), len(mask)) + start2 = min(max(start2, 0), len(mask)) + stop1 = min(max(stop1, 0), len(mask)) + stop2 = min(max(stop2, 0), len(mask)) + + if stop1 >= start2: + found = True + merged.append((start1, stop2)) + index += 2 + else: + merged.append((start1, stop1)) + index += 1 + if index == len(ranges) - 1: + merged.append((start2, stop2)) + ranges = merged + if not found: + for index in range(1, len(ranges)): + start1, stop1 = ranges[index - 1] + start2, stop2 = ranges[index] + assert start1 < stop1 + assert start2 < stop2 + assert start1 < start2 + assert stop1 < stop2 + assert stop1 < start2 break - if segments is None: - canvas = img.copy() - else: - canvas = np.hstack(segments) + segments = [] + starts_ = [] + stops_ = [] + domain = img.shape[1] + widths = [] + total_width = 0 + for start, stop in ranges: + segment = img[:, start:stop] + segments.append(segment) + + starts_.append(int(round(spectrogram.duration * (start / domain)))) + stops_.append(int(round(spectrogram.duration * (stop / domain)))) + widths.append(stop - start) + total_width += stop - start + + # buffer = np.zeros((len(img), 20, 3), dtype=img.dtype) + # segments.append(buffer) + # segments = segments[:-1] + + if len(segments) > 0: + break + + threshold -= 0.05 + if threshold < 0: + segments = None + break + + if segments is None: + canvas = img.copy() + else: + canvas = np.hstack(segments) - canvas = Image.fromarray(canvas, 'RGB') + canvas = Image.fromarray(canvas, 'RGB') - canvas.save('temp.compressed.jpg') + canvas.save('temp.compressed.jpg') - buf = io.BytesIO() - canvas.save(buf, format='JPEG', quality=80) - buf.seek(0) - name = f'{spectrogram.pk}_spectrogram_compressed.jpg' - image_file = File(buf, name=name) + buf = io.BytesIO() + canvas.save(buf, format='JPEG', quality=80) + buf.seek(0) + name = f'{spectrogram.pk}_spectrogram_compressed.jpg' + image_file = File(buf, name=name) - return total_width, image_file, widths, starts, stops + return total_width, image_file, widths, starts, stops @shared_task @@ -200,7 +197,5 @@ def generate_compress_spectrogram(recording_id: int, spectrogram_id: int): widths=widths, starts=starts, stops=stops, - cache_invalidated=False + cache_invalidated=False, ) - - From 1c38d150c075924a69f4f373b26f8c27370e8c83 Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Wed, 17 Apr 2024 10:30:42 -0400 Subject: [PATCH 6/8] change spectrogram to utilzie URLs instead of base64 --- bats_ai/core/models/compressed_spectrogram.py | 16 ++-- bats_ai/core/models/spectrogram.py | 82 ++----------------- bats_ai/core/tasks.py | 25 ++++-- bats_ai/core/views/recording.py | 44 +++++----- client/src/components/ThumbnailViewer.vue | 2 +- client/src/views/Spectrogram.vue | 12 ++- 6 files changed, 73 insertions(+), 108 deletions(-) diff --git a/bats_ai/core/models/compressed_spectrogram.py b/bats_ai/core/models/compressed_spectrogram.py index 04466bf..5dd57c1 100644 --- a/bats_ai/core/models/compressed_spectrogram.py +++ b/bats_ai/core/models/compressed_spectrogram.py @@ -1,12 +1,14 @@ from PIL import Image import cv2 from django.contrib.postgres.fields import ArrayField +from django.core.files.storage import default_storage from django.db import models from django.dispatch import receiver from django_extensions.db.models import TimeStampedModel import numpy as np -from bats_ai.core.models import Recording, Spectrogram +from .recording import Recording +from .spectrogram import Spectrogram FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -18,12 +20,16 @@ class CompressedSpectrogram(TimeStampedModel, models.Model): recording = models.ForeignKey(Recording, on_delete=models.CASCADE) spectrogram = models.ForeignKey(Spectrogram, on_delete=models.CASCADE) image_file = models.FileField() - length = models.IntegerField() # pixels - starts = ArrayField(ArrayField(models.IntegerField())) # pixels - stops = ArrayField(ArrayField(models.IntegerField())) # milliseconds - widths = ArrayField(ArrayField(models.IntegerField())) # hz + length = models.IntegerField() + starts = ArrayField(ArrayField(models.IntegerField())) + stops = ArrayField(ArrayField(models.IntegerField())) + widths = ArrayField(ArrayField(models.IntegerField())) cache_invalidated = models.BooleanField(default=True) + @property + def image_url(self): + return default_storage.url(self.image_file.name) + def predict(self): import json import os diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index a31ffee..c4b38cb 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -1,10 +1,12 @@ import base64 import io import math +import os from PIL import Image import cv2 from django.core.files import File +from django.core.files.storage import default_storage from django.db import models from django.db.models.fields.files import FieldFile from django.dispatch import receiver @@ -14,7 +16,7 @@ import numpy as np import tqdm -from bats_ai.core.models import Recording +from .recording import Recording FREQ_MIN = 5e3 FREQ_MAX = 120e3 @@ -201,6 +203,7 @@ def generate(cls, recording, colormap=None, dpi=600): colormap=colormap, ) spectrogram.save() + os.remove(f'/tmp/temp.{colormap}.jpg') return spectrogram.pk @property @@ -223,80 +226,9 @@ def base64(self): return img_base64 - def predict(self): - import json - import os - - import onnx - import onnxruntime as ort - import tqdm - - img, _, _ = self.compressed - - relative = ('..',) * 4 - asset_path = os.path.abspath(os.path.join(__file__, *relative, 'assets')) - - onnx_filename = os.path.join(asset_path, 'model.mobilenet.onnx') - assert os.path.exists(onnx_filename) - - session = ort.InferenceSession( - onnx_filename, - providers=[ - ( - 'CUDAExecutionProvider', - { - 'cudnn_conv_use_max_workspace': '1', - 'device_id': 0, - 'cudnn_conv_algo_search': 'HEURISTIC', - }, - ), - 'CPUExecutionProvider', - ], - ) - - img = np.array(img) - - h, w, c = img.shape - ratio_y = 224 / h - ratio_x = ratio_y * 0.5 - raw = cv2.resize(img, None, fx=ratio_x, fy=ratio_y, interpolation=cv2.INTER_LANCZOS4) - - h, w, c = raw.shape - if w <= h: - canvas = np.zeros((h, h + 1, 3), dtype=raw.dtype) - canvas[:, :w, :] = raw - raw = canvas - h, w, c = raw.shape - - inputs_ = [] - for index in range(0, w - h, 100): - inputs_.append(raw[:, index : index + h, :]) - inputs_.append(raw[:, -h:, :]) - inputs_ = np.array(inputs_) - - chunksize = 1 - chunks = np.array_split(inputs_, np.arange(chunksize, len(inputs_), chunksize)) - outputs = [] - for chunk in tqdm.tqdm(chunks, desc='Inference'): - outputs_ = session.run( - None, - {'input': chunk}, - ) - outputs.append(outputs_[0]) - outputs = np.vstack(outputs) - outputs = outputs.mean(axis=0) - - model = onnx.load(onnx_filename) - mapping = json.loads(model.metadata_props[0].value) - labels = [mapping['forward'][str(index)] for index in range(len(mapping['forward']))] - - prediction = np.argmax(outputs) - label = labels[prediction] - score = outputs[prediction] - - confs = dict(zip(labels, outputs)) - - return label, score, confs + @property + def image_url(self): + return default_storage.url(self.image_file.name) @receiver(models.signals.pre_delete, sender=Spectrogram) diff --git a/bats_ai/core/tasks.py b/bats_ai/core/tasks.py index b1abd90..a30477b 100644 --- a/bats_ai/core/tasks.py +++ b/bats_ai/core/tasks.py @@ -1,4 +1,5 @@ import io +import tempfile from PIL import Image from celery import shared_task @@ -11,7 +12,6 @@ def generate_compressed(recording: Recording, spectrogram: Spectrogram): - print(spectrogram) img = spectrogram.image_np annotations = Annotations.objects.filter(recording=recording) @@ -137,15 +137,21 @@ def generate_compressed(recording: Recording, spectrogram: Spectrogram): canvas = Image.fromarray(canvas, 'RGB') - canvas.save('temp.compressed.jpg') + # Use temporary files + with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file: + temp_file_name = temp_file.name + canvas.save(temp_file_name) - buf = io.BytesIO() - canvas.save(buf, format='JPEG', quality=80) - buf.seek(0) + # Read the temporary file + with open(temp_file_name, 'rb') as f: + temp_file_content = f.read() + + # Wrap the content in BytesIO + buf = io.BytesIO(temp_file_content) name = f'{spectrogram.pk}_spectrogram_compressed.jpg' image_file = File(buf, name=name) - return total_width, image_file, widths, starts, stops + return total_width, image_file, widths, starts_, stops_ @shared_task @@ -199,3 +205,10 @@ def generate_compress_spectrogram(recording_id: int, spectrogram_id: int): stops=stops, cache_invalidated=False, ) + + +@shared_task +def predict(compressed_spectrogram_id: int): + compressed_spectrogram = CompressedSpectrogram.objects.get(pk=compressed_spectrogram_id) + label, score, confs = compressed_spectrogram.predict() + return label, score, confs diff --git a/bats_ai/core/views/recording.py b/bats_ai/core/views/recording.py index 17be8df..44321ce 100644 --- a/bats_ai/core/views/recording.py +++ b/bats_ai/core/views/recording.py @@ -11,7 +11,14 @@ from ninja.files import UploadedFile from ninja.pagination import RouterPaginated -from bats_ai.core.models import Annotations, Recording, Species, TemporalAnnotations, colormap +from bats_ai.core.models import ( + Annotations, + CompressedSpectrogram, + Recording, + Species, + TemporalAnnotations, + colormap, +) from bats_ai.core.tasks import recording_compute_spectrogram from bats_ai.core.views.species import SpeciesSchema from bats_ai.core.views.temporal_annotations import ( @@ -272,7 +279,7 @@ def get_spectrogram(request: HttpRequest, id: int): spectrogram = recording.spectrogram spectro_data = { - 'base64_spectrogram': spectrogram.base64, + 'url': spectrogram.image_url, 'spectroInfo': { 'width': spectrogram.width, 'height': spectrogram.height, @@ -328,30 +335,29 @@ def get_spectrogram(request: HttpRequest, id: int): def get_spectrogram_compressed(request: HttpRequest, id: int): try: recording = Recording.objects.get(pk=id) - except Recording.DoesNotExist: - return {'error': 'Recording not found'} + compressed_spectrogram = CompressedSpectrogram.objects.filter(recording=id).first() + except compressed_spectrogram.DoesNotExist: + return {'error': 'Compressed Spectrogram'} + except recording.DoesNotExist: + return {'error': 'Recording does not exist'} with colormap(): - label, score, confs = recording.spectrogram.predict() + label, score, confs = compressed_spectrogram.predict() print(label, score, confs) - with colormap(None): - spectrogram = recording.spectrogram - _, compressed_base64, metadata = spectrogram.compressed - spectro_data = { - 'base64_spectrogram': compressed_base64, + 'url': compressed_spectrogram.image_url, 'spectroInfo': { - 'width': spectrogram.width, + 'width': compressed_spectrogram.spectrogram.width, 'start_time': 0, - 'end_time': spectrogram.duration, - 'height': spectrogram.height, - 'low_freq': spectrogram.frequency_min, - 'high_freq': spectrogram.frequency_max, - 'start_times': metadata['starts'], - 'end_times': metadata['stops'], - 'widths': metadata['widths'], - 'compressedWidth': metadata['length'], + 'end_time': compressed_spectrogram.spectrogram.duration, + 'height': compressed_spectrogram.spectrogram.height, + 'low_freq': compressed_spectrogram.spectrogram.frequency_min, + 'high_freq': compressed_spectrogram.spectrogram.frequency_max, + 'start_times': compressed_spectrogram.starts, + 'end_times': compressed_spectrogram.stops, + 'widths': compressed_spectrogram.widths, + 'compressedWidth': compressed_spectrogram.length, }, } diff --git a/client/src/components/ThumbnailViewer.vue b/client/src/components/ThumbnailViewer.vue index 3ba31f5..0904aee 100644 --- a/client/src/components/ThumbnailViewer.vue +++ b/client/src/components/ThumbnailViewer.vue @@ -160,7 +160,7 @@ export default defineComponent({ clientHeight.value = containerRef.value.clientHeight; } if (containerRef.value && ! geoJS.getGeoViewer().value) { - geoJS.initializeViewer(containerRef.value, naturalWidth, naturalHeight, true); + geoJS.initializeViewer(containerRef.value, naturalWidth, naturalHeight, true); } const coords = geoJS.getGeoViewer().value.camera().worldToDisplay({x: 0, y:0}); const end = geoJS.getGeoViewer().value.camera().worldToDisplay({x: 0, y:naturalHeight}); diff --git a/client/src/views/Spectrogram.vue b/client/src/views/Spectrogram.vue index 2147323..81819a0 100644 --- a/client/src/views/Spectrogram.vue +++ b/client/src/views/Spectrogram.vue @@ -84,17 +84,25 @@ export default defineComponent({ }; const loadData = async () => { + loadedImage.value = false; const response = compressed.value ? await getSpectrogramCompressed(props.id) : await getSpectrogram(props.id); - image.value.src = `data:image/png;base64,${response.data["base64_spectrogram"]}`; + if (response.data['url']) { + image.value.src = response.data['url']; + } else { + // TODO Error Out if there is no URL + console.error('No URL found for the spectrogram'); + } + image.value.onload = () => { + loadedImage.value = true; + }; spectroInfo.value = response.data["spectroInfo"]; annotations.value = response.data["annotations"]?.sort((a, b) => a.start_time - b.start_time) || []; temporalAnnotations.value = response.data["temporal"]?.sort((a, b) => a.start_time - b.start_time) || []; if (response.data.currentUser) { currentUser.value = response.data.currentUser; } - loadedImage.value = true; const speciesResponse = await getSpecies(); speciesList.value = speciesResponse.data; if (response.data.otherUsers && spectroInfo.value) { From 6653d06f2802fd65c7c644503ffd9cfa79faa41e Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Wed, 17 Apr 2024 11:14:37 -0400 Subject: [PATCH 7/8] update viewer defaults --- client/src/components/ThumbnailViewer.vue | 4 ++-- client/src/components/geoJS/geoJSUtils.ts | 22 +++++++++++++++------- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/client/src/components/ThumbnailViewer.vue b/client/src/components/ThumbnailViewer.vue index 0904aee..d766862 100644 --- a/client/src/components/ThumbnailViewer.vue +++ b/client/src/components/ThumbnailViewer.vue @@ -235,8 +235,8 @@ export default defineComponent({ position: absolute; top: 50%; left: 50%; - -ms-transform: translate(-50%, -50%); - transform: translate(-50%, -50%); + // -ms-transform: translate(-50%, -50%); + // transform: translate(-50%, -50%); } .geojs-map.annotation-input { cursor: inherit; diff --git a/client/src/components/geoJS/geoJSUtils.ts b/client/src/components/geoJS/geoJSUtils.ts index c799cba..0269de5 100644 --- a/client/src/components/geoJS/geoJSUtils.ts +++ b/client/src/components/geoJS/geoJSUtils.ts @@ -17,6 +17,8 @@ const useGeoJS = () => { right: 1, }; + let originalDimensions = { width: 0, height: 0 }; + const getGeoViewer = () => { return geoViewer; }; @@ -29,6 +31,7 @@ const useGeoJS = () => { ) => { thumbnail.value = thumbanilVal; container.value = sourceContainer; + originalDimensions = {width, height }; const params = geo.util.pixelCoordinateParams(container.value, width, height); if (!container.value) { return; @@ -112,7 +115,7 @@ const useGeoJS = () => { .draw(); } if (resetCam) { - resetMapDimensions(width, height, 0.3,resetCam); + resetMapDimensions(width, height, 0.3, resetCam); } else { const params = geo.util.pixelCoordinateParams(container.value, width, height, width, height); const margin = 0.3; @@ -129,16 +132,21 @@ const useGeoJS = () => { }; const resetZoom = () => { - const { width: mapWidth } = geoViewer.value.camera().viewport; + const { width: mapWidth, } = geoViewer.value.camera().viewport; const bounds = !thumbnail.value ? { - left: -125, // Making sure the legend is on the screen - top: 0, - right: mapWidth, + left: 0, // Making sure the legend is on the screen + top: -(originalBounds.bottom - originalDimensions.height) / 2.0, + right: mapWidth*2, bottom: originalBounds.bottom, } - : originalBounds; + : { + left: 0, + top: 0, + right: originalDimensions.width, + bottom: originalDimensions.height, + }; const zoomAndCenter = geoViewer.value.zoomAndCenterFromBounds(bounds, 0); geoViewer.value.zoom(zoomAndCenter.zoom); geoViewer.value.center(zoomAndCenter.center); @@ -154,13 +162,13 @@ const useGeoJS = () => { }); const params = geo.util.pixelCoordinateParams(container.value, width, height, width, height); const { right, bottom } = params.map.maxBounds; - originalBounds = params.map.maxBounds; geoViewer.value.maxBounds({ left: 0 - right * margin, top: 0 - bottom * margin, right: right * (1 + margin), bottom: bottom * (1 + margin), }); + originalBounds = geoViewer.value.maxBounds(); geoViewer.value.zoomRange({ // do not set a min limit so that bounds clamping determines min min: -Infinity, From 6adb2052db825ddb7e5ca488cdfac13f6958ec7a Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Fri, 19 Apr 2024 10:05:34 -0400 Subject: [PATCH 8/8] setup update, pip upgrade, merged testing --- .../migrations/0010_compressedspectrogram.py | 2 +- .../migrations/0011_spectrogram_colormap.py | 17 -------- bats_ai/core/models/compressed_spectrogram.py | 1 - bats_ai/core/models/spectrogram.py | 7 +--- bats_ai/core/rest/__init__.py | 4 +- bats_ai/core/rest/image.py | 42 ------------------- dev/django.Dockerfile | 2 + setup.py | 10 ++--- 8 files changed, 10 insertions(+), 75 deletions(-) delete mode 100644 bats_ai/core/migrations/0011_spectrogram_colormap.py delete mode 100644 bats_ai/core/rest/image.py diff --git a/bats_ai/core/migrations/0010_compressedspectrogram.py b/bats_ai/core/migrations/0010_compressedspectrogram.py index 0ce5b18..28fe4fb 100644 --- a/bats_ai/core/migrations/0010_compressedspectrogram.py +++ b/bats_ai/core/migrations/0010_compressedspectrogram.py @@ -1,4 +1,4 @@ -# Generated by Django 4.1.13 on 2024-04-12 18:56 +# Generated by Django 4.1.13 on 2024-04-19 13:55 import django.contrib.postgres.fields from django.db import migrations, models diff --git a/bats_ai/core/migrations/0011_spectrogram_colormap.py b/bats_ai/core/migrations/0011_spectrogram_colormap.py deleted file mode 100644 index bc22fd8..0000000 --- a/bats_ai/core/migrations/0011_spectrogram_colormap.py +++ /dev/null @@ -1,17 +0,0 @@ -# Generated by Django 4.1.13 on 2024-04-09 13:49 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ('core', '0010_recording_computed_species_recording_detector_and_more'), - ] - - operations = [ - migrations.AddField( - model_name='spectrogram', - name='colormap', - field=models.CharField(max_length=20, null=True), - ), - ] diff --git a/bats_ai/core/models/compressed_spectrogram.py b/bats_ai/core/models/compressed_spectrogram.py index 3717c18..5dd57c1 100644 --- a/bats_ai/core/models/compressed_spectrogram.py +++ b/bats_ai/core/models/compressed_spectrogram.py @@ -40,7 +40,6 @@ def predict(self): img = Image.open(self.image_file) - relative = ('..',) * 4 asset_path = os.path.abspath(os.path.join(__file__, *relative, 'assets')) diff --git a/bats_ai/core/models/spectrogram.py b/bats_ai/core/models/spectrogram.py index 134a56c..eaedefb 100644 --- a/bats_ai/core/models/spectrogram.py +++ b/bats_ai/core/models/spectrogram.py @@ -2,7 +2,6 @@ import io import logging import math -import os from PIL import Image import cv2 @@ -15,12 +14,9 @@ import librosa import matplotlib.pyplot as plt import numpy as np - -from .recording import Recording -import scipy import tqdm -from bats_ai.core.models import Annotations, Recording +from .recording import Recording logger = logging.getLogger(__name__) @@ -237,7 +233,6 @@ def generate(cls, recording, colormap=None, dpi=520): colormap=colormap, ) spectrogram.save() - os.remove(f'/tmp/temp.{colormap}.jpg') return spectrogram.pk @property diff --git a/bats_ai/core/rest/__init__.py b/bats_ai/core/rest/__init__.py index ecd2d00..73bcd90 100644 --- a/bats_ai/core/rest/__init__.py +++ b/bats_ai/core/rest/__init__.py @@ -1,10 +1,8 @@ from rest_framework import routers -from .image import ImageViewSet from .spectrogram import SpectrogramViewSet -__all__ = ['ImageViewSet', 'SpectrogramViewSet'] +__all__ = ['SpectrogramViewSet'] rest = routers.SimpleRouter() -rest.register(r'images', ImageViewSet) rest.register(r'spectrograms', SpectrogramViewSet) diff --git a/bats_ai/core/rest/image.py b/bats_ai/core/rest/image.py deleted file mode 100644 index d749110..0000000 --- a/bats_ai/core/rest/image.py +++ /dev/null @@ -1,42 +0,0 @@ -from django.http import HttpResponseRedirect -from django_filters import rest_framework as filters -from rest_framework import serializers, status -from rest_framework.decorators import action -from rest_framework.pagination import PageNumberPagination -from rest_framework.permissions import IsAuthenticatedOrReadOnly -from rest_framework.response import Response -from rest_framework.viewsets import ReadOnlyModelViewSet - -from bats_ai.core.models import Image -from bats_ai.core.tasks import image_compute_checksum - - -class ImageSerializer(serializers.ModelSerializer): - class Meta: - model = Image - fields = ['id', 'name', 'checksum', 'created', 'owner'] - read_only_fields = ['checksum', 'created'] - - -class ImageViewSet(ReadOnlyModelViewSet): - queryset = Image.objects.all() - - permission_classes = [IsAuthenticatedOrReadOnly] - serializer_class = ImageSerializer - - filter_backends = [filters.DjangoFilterBackend] - filterset_fields = ['name', 'checksum'] - - pagination_class = PageNumberPagination - - @action(detail=True, methods=['get']) - def download(self, request, pk=None): - image = self.get_object() - return HttpResponseRedirect(image.blob.url) - - @action(detail=True, methods=['post']) - def compute(self, request, pk=None): - # Ensure that the image exists, so a non-existent pk isn't dispatched - image = self.get_object() - image_compute_checksum.delay(image.pk) - return Response('', status=status.HTTP_202_ACCEPTED) diff --git a/dev/django.Dockerfile b/dev/django.Dockerfile index 5992260..7861630 100644 --- a/dev/django.Dockerfile +++ b/dev/django.Dockerfile @@ -28,6 +28,8 @@ COPY ./setup.py /opt/django-project/setup.py # Use a directory name which will never be an import name, as isort considers this as first-party. WORKDIR /opt/django-project +# hadolint ignore=DL3013 +RUN pip install --no-cache-dir --upgrade pip RUN set -ex \ && pip install --no-cache-dir -e .[dev] diff --git a/setup.py b/setup.py index 2602802..02f4b0f 100644 --- a/setup.py +++ b/setup.py @@ -44,7 +44,6 @@ 'django-allauth', 'django-configurations[database,email]', 'django-extensions', - 'django-filter', 'django-large-image', 'django-oauth-toolkit', 'djangorestframework', @@ -55,20 +54,21 @@ 'django-s3-file-field[boto3]<1', 'gunicorn', 'flower', - 'large-image[rasterio,pil]>=1.22', + # Spectrogram Generation 'librosa', 'matplotlib', 'mercantile', 'numpy', + # 'onnxruntime-gpu', 'onnx', 'onnxruntime', - # 'onnxruntime-gpu', 'opencv-python-headless', 'tqdm', - 'django-large-image', + # large image + 'django-large-image>=0.10.0', 'large-image[rasterio,pil]>=1.22', 'rio-cogeo', - 'mercantile', + # guano metadata 'guano', ], extras_require={