sound-annotator/spec_viewer.py at master · conservationtechlab/sound-annotator · 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
from PyQt5 import QtCore
from PyQt5.QtWidgets import QLabel
from PyQt5.QtGui import QPainter, QBrush, QPen, QImage, QPixmap, QColor
import logging
from math import ceil

from audio_clip import AudioClip


class SpecViewer(QLabel):

    show_new_file_info = QtCore.pyqtSignal(logging.Logger)

    def __init__(self, parent):
        super().__init__(parent=parent)
        self.logger = logging.getLogger(self.__class__.__name__)
        self.c1 = (0, 0)
        self.c2 = (0, 0)
        self.curr_clip = None
        self.min_dur = None
        self.pen_weight = 1
        self.h = self.pen_weight

    def get_curr_selection(self):
        start_spec = min(self.c1[0], self.c2[0])  # get start and end in spectrogram space
        end_spec = max(self.c1[0], self.c2[0])

        _, spec_w = self.curr_clip.spec.shape
        start = int(start_spec * len(self.curr_clip.data) / spec_w)
        end   = int(end_spec   * len(self.curr_clip.data) / spec_w)
        return start, end

    def mousePressEvent(self, e):
        self.logger.debug('mouse pressed at ({}, {})'.format(e.x(), e.y()))
        self.c1 = e.x(), e.y()
        self.c2 = e.x(), e.y()
        self.update()

    def mouseMoveEvent(self, e):
        self.c2 = e.x(), e.y()
        self.update()

    def mouseReleaseEvent(self, e):
        self.logger.debug('mouse released at ({}, {})'.format(e.x(), e.y()))
        self.c2 = e.x(), e.y()
        self.update()
        self.save_selection()

    def paintEvent(self, e):
        super(SpecViewer, self).paintEvent(e)
        painter = QPainter(self)
        x, _ = self.c1
        w = self.c2[0] - self.c1[0]
        painter.setPen(QPen(QColor(178, 217, 181), self.pen_weight, QtCore.Qt.SolidLine))
        painter.setBrush(QBrush(QColor(178, 217, 181), QtCore.Qt.Dense7Pattern))
        painter.drawRect(x, self.pen_weight, w, self.h - self.pen_weight)

    def new_clip(self, fp):
        self.c1 = (0, 0)
        self.c2 = (0, 0)
        self.update()
        self.curr_clip = AudioClip(fp)
        viewable_spec = (self.curr_clip.spec * 255).astype('uint8')
        self.h, w = viewable_spec.shape
        self.setPixmap(QPixmap.fromImage(QImage(viewable_spec, w, self.h,
            int(viewable_spec.nbytes / self.h), QImage.Format_Grayscale8)))
        self.show_new_file_info.emit(self.logger)

    def save_selection(self):
        start_spec = min(self.c1[0], self.c2[0])  # get start and end in spectrogram space
        end_spec = max(self.c1[0], self.c2[0])

        _, spec_w = self.curr_clip.spec.shape  # get min dur in spec space
        min_dur_spec = self.min_dur * self.curr_clip.sr * spec_w / len(self.curr_clip.data)

        dur_spec = end_spec - start_spec
        if dur_spec < min_dur_spec:
            pad_len = ceil((min_dur_spec - dur_spec) / 2)
            start_spec = max(start_spec - pad_len, 0)
            end_spec = min(end_spec + pad_len, spec_w)

            self.c1 = (start_spec, self.c1[1])
            self.c2 = (end_spec, self.c2[1])
            self.update()

        self.curr_clip.write_mini_clip(start_spec, end_spec)