diff --git a/get_spectrogram.py b/get_spectrogram.py index a058589eb855b34bf05b98b50123fec579cf0c91..8ed5d3501c33a803899fcada1ae637eb5b562a71 100755 --- a/get_spectrogram.py +++ b/get_spectrogram.py @@ -54,7 +54,7 @@ def main(data, arguments): utils.create_spectrogram( sig, arguments.directory, name, arguments.cmap, window_size=int(arguments.window), - overlap=arguments.hop) + overlap=arguments.hop, vmin=arguments.vmin) except Exception as error: folder = 'spectrograms' @@ -85,6 +85,8 @@ if __name__ == "__main__": help='Overlap in secondes between 2 spectrograms', default=0) parser.add_argument('--rf', type=int, help='Resampling Frequency of the signal. If no argument,' ' will be original frequency sampling of the recording', default=None) + parser.add_argument('--vmin', type=str, help="If vmin == True, then the spectrogram's minimum color" + ' will be stft.mean(). If False stft.min()', default=True) parser.add_argument( '--cpu', type=int, help='To speed up the process, write 2 or more', default=1) parser.add_argument( diff --git a/utils.py b/utils.py index 29550eaecad2961e605657a8b9f3ace12b100927..f55344a85f562525f7136e0b2262f68da94e5a0a 100755 --- a/utils.py +++ b/utils.py @@ -87,7 +87,7 @@ def signal_processing(sig, rf, fs, high=None, low=None): return sig -def create_spectrogram(sig, directory, names, cmap, window_size=1024, overlap=.5,): +def create_spectrogram(sig, directory, names, cmap, window_size=1024, overlap=.5, minimum=True): """ Create a spectrogram STFT with hanning window and save it into a directory @@ -107,9 +107,13 @@ def create_spectrogram(sig, directory, names, cmap, window_size=1024, overlap=.5 hop_length=int(overlap_size), window='hann') # Compute the STFT stft = np.log10(np.abs(stft)) # Adapt the Complex-valued matrix fig = plt.figure() + if minimum: + vmin = stft.mean() + else: + vmin = stft.min() # plot the spectrogram plt.imshow(stft[::-1], aspect='auto', - interpolation=None, cmap=cmap) # you can add : vmin=stft.mean()) + interpolation=None, cmap=cmap) # you can add : vmin=vmin) # Remove all the borders around the plot plt.subplots_adjust(top=1, bottom=0, left=0, right=1) if names: diff --git a/yolov5/detect.py b/yolov5/detect.py index b035c739cc9b6a2c160470033a4c31e249ddb3f6..1b6c7a52db62b3c302d9ebe2b9a278999d9fce8e 100755 --- a/yolov5/detect.py +++ b/yolov5/detect.py @@ -74,6 +74,7 @@ def run( rf=22050, window=1024, hop=0.5, + minimum=True, # vmin spectrogram is stft.mean() low=None, high=None, cmap='viridis', @@ -136,7 +137,7 @@ def run( print(f'You put hop > 1, this has been corrected by putting hop at {hop}') else: hop = window * hop - dataset = LoadSpectros(source, sampleDur, rf, window, hop, low, high, cmap, img_size=imgsz, stride=stride, auto=pt) + dataset = LoadSpectros(source, sampleDur, rf, window, hop, low, high, cmap, img_size=imgsz, stride=stride, auto=pt, minimum=minimum) else: dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) vid_path, vid_writer = [None] * bs, [None] * bs @@ -279,6 +280,8 @@ def parse_opt(): help="Colormap used for the spectrograms",type=str) parser.add_argument('--window', default=1024, help="Window size for each spectrogram for detection",type=int) parser.add_argument('--hop', default=0.5, help="Hop lenght for each spectrogram for detection",type=float) + parser.add_argument('--vmin', type=str, help='If vmin == True, then the spectrogram minimum color' + ' will be stft.mean(). If False stft.min()', default=True) parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences') parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') parser.add_argument('--sound', default=False, action='store_true') diff --git a/yolov5/utils/dataloaders.py b/yolov5/utils/dataloaders.py index 5c351a82c068155edbbed51fb63a657a1cb1110e..cbc3738bc75b014636a5c695fda40394cfedbd75 100755 --- a/yolov5/utils/dataloaders.py +++ b/yolov5/utils/dataloaders.py @@ -239,8 +239,8 @@ class LoadScreenshots: return str(self.screen), im, im0, None, s # screen, img, original img, im0s, s class LoadSpectros: - def __init__(self, folder, sampleDur, rf, window, hop, low, high, cmap, img_size, stride=32, auto=True): - self.folder, self.sampleDur, self.rf, self.window, self.hop, self.low, self.high, self.cmap, self.img_size, self.stride, self.auto = folder, sampleDur, rf, window, hop, low, high, cmap, img_size, stride, auto + def __init__(self, folder, sampleDur, rf, window, hop, low, high, cmap, img_size, stride=32, auto=True, minimum=True): + self.folder, self.sampleDur, self.rf, self.window, self.hop, self.low, self.high, self.cmap, self.img_size, self.stride, self.auto, self.minimum = folder, sampleDur, rf, window, hop, low, high, cmap, img_size, stride, auto, minimum self.files = os.listdir(folder) self.mode = 'image' self.samples = [] @@ -290,8 +290,12 @@ class LoadSpectros: stft = librosa.stft(sig, n_fft=self.window, hop_length=hop, window='hann') # Compute the STFT stft = np.log10(np.abs(stft)) + if minimum: + vmin = stft.mean() + else: + vmin = stft.min() axim = plt.imshow(stft, aspect = "auto", interpolation = None, - cmap = self.cmap) # you can add : vmin=np.mean(stft)) + cmap = self.cmap, vmin=vmin) plt.subplots_adjust(top=1, bottom=0, left=0, right=1) im0 = axim.make_image(fig.canvas)[0][:,:,:-1][:,:,::-1] cv2.imwrite(path, im0)