batai/bats_ai/utils/spectrogram_utils.py at 3822131f351b2ecd861a607ad4b291a54bc34cb5 · Kitware/batai · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
from __future__ import annotations

import logging
import os
from typing import NotRequired, TypedDict

from django.contrib.contenttypes.models import ContentType
from django.core.files import File

from bats_ai.core.models import SpectrogramImage
from bats_ai.core.models.nabat import NABatCompressedSpectrogram, NABatRecording, NABatSpectrogram
from bats_ai.core.utils.image_utils import waveplot_to_grayscale_transparent

logger = logging.getLogger(__name__)


class SpectrogramAssetResult(TypedDict):
    paths: list[str]
    waveplot_paths: NotRequired[list[str]]
    width: int
    height: int


class SpectrogramCompressedAssetResult(TypedDict):
    paths: list[str]
    masks: list[str]
    waveplot_paths: NotRequired[list[str]]
    width: int
    height: int
    widths: list[float]
    starts: list[float]
    stops: list[float]


class SpectrogramAssets(TypedDict):
    duration: float
    freq_min: int
    freq_max: int
    normal: SpectrogramAssetResult
    compressed: SpectrogramCompressedAssetResult


class PredictionOutput(TypedDict):
    label: str
    score: float
    confs: dict[str, float]


def generate_nabat_spectrogram(
    nabat_recording: NABatRecording, results: SpectrogramAssets
) -> NABatSpectrogram:
    spectrogram, _ = NABatSpectrogram.objects.get_or_create(
        nabat_recording=nabat_recording,
        defaults={
            "width": results["normal"]["width"],
            "height": results["normal"]["height"],
            "duration": results["duration"],
            "frequency_min": results["freq_min"],
            "frequency_max": results["freq_max"],
        },
    )

    # Create SpectrogramImage objects for each normal image
    for idx, img_path in enumerate(results["normal"]["paths"]):
        with open(img_path, "rb") as f:
            SpectrogramImage.objects.get_or_create(
                content_type=ContentType.objects.get_for_model(spectrogram),
                object_id=spectrogram.id,
                index=idx,
                defaults={
                    "image_file": File(f, name=os.path.basename(img_path)),
                    "type": "spectrogram",
                },
            )

    for idx, img_path in enumerate(results["normal"].get("waveplot_paths", [])):
        buf = waveplot_to_grayscale_transparent(img_path)
        base = os.path.splitext(os.path.basename(img_path))[0]
        SpectrogramImage.objects.get_or_create(
            content_type=ContentType.objects.get_for_model(spectrogram),
            object_id=spectrogram.id,
            index=idx,
            defaults={
                "image_file": File(buf, name=f"{base}.png"),
                "type": "waveform_uncompressed",
            },
        )

    return spectrogram


def generate_nabat_compressed_spectrogram(
    nabat_recording: NABatRecording,
    spectrogram: NABatSpectrogram,
    compressed_results: SpectrogramCompressedAssetResult,
) -> NABatCompressedSpectrogram:
    compressed_obj, _ = NABatCompressedSpectrogram.objects.get_or_create(
        nabat_recording=nabat_recording,
        spectrogram=spectrogram,
        defaults={
            "length": compressed_results["width"],
            "widths": compressed_results["widths"],
            "starts": compressed_results["starts"],
            "stops": compressed_results["stops"],
            "cache_invalidated": False,
        },
    )

    # Save compressed images
    for idx, img_path in enumerate(compressed_results["paths"]):
        with open(img_path, "rb") as f:
            SpectrogramImage.objects.get_or_create(
                content_type=ContentType.objects.get_for_model(compressed_obj),
                object_id=compressed_obj.id,
                index=idx,
                defaults={
                    "image_file": File(f, name=os.path.basename(img_path)),
                    "type": "compressed",
                },
            )

    # Save mask images (from batbot metadata mask_path)
    for idx, mask_path in enumerate(compressed_results.get("masks", [])):
        with open(mask_path, "rb") as f:
            SpectrogramImage.objects.get_or_create(
                content_type=ContentType.objects.get_for_model(compressed_obj),
                object_id=compressed_obj.id,
                index=idx,
                type="masks",
                defaults={
                    "image_file": File(f, name=os.path.basename(mask_path)),
                },
            )

    for idx, waveplot_path in enumerate(compressed_results.get("waveplot_paths", [])):
        buf = waveplot_to_grayscale_transparent(waveplot_path)
        base = os.path.splitext(os.path.basename(waveplot_path))[0]
        SpectrogramImage.objects.get_or_create(
            content_type=ContentType.objects.get_for_model(compressed_obj),
            object_id=compressed_obj.id,
            index=idx,
            type="waveform_compressed",
            defaults={
                "image_file": File(buf, name=f"{base}.png"),
            },
        )

    return compressed_obj