import multiprocessing import sys import traceback from itertools import groupby from time import time from typing import Dict, List, Tuple import dejavu.logic.decoder as decoder from dejavu.base_classes.base_database import get_database from dejavu.config.settings import (DEFAULT_FS, DEFAULT_OVERLAP_RATIO, DEFAULT_WINDOW_SIZE, FIELD_FILE_SHA1, FIELD_SONGNAME, FIELD_TOTAL_HASHES, FIELD_EXTENSION, FINGERPRINTED_CONFIDENCE, FINGERPRINTED_HASHES, HASHES_MATCHED, INPUT_CONFIDENCE, INPUT_HASHES, OFFSET, OFFSET_SECS, SONG_ID, SONG_NAME, TOPN) from dejavu.logic.fingerprint import fingerprint from s3_file import S3File from config import client, s3 class Dejavu: def __init__(self, config): self.config = config # initialize db db_cls = get_database(config.get('database_type', 'mysql').lower()) self.db = db_cls(**config.get('database', {})) self.db.setup() # if we should limit seconds fingerprinted, # None|-1 means use entire track self.limit = self.config.get("fingerprint_limit", None) if self.limit == -1: # for JSON compatibility self.limit = None self.songs = None self.filenames = None def load_fingerprinted_filenames(self) -> None: """ Keeps a dictionary with the hashes of the fingerprinted songs, in that way is possible to check whether or not an audio file was already processed. """ # get songs previously indexed self.songs = self.db.get_songs() self.filenames = set() for song in self.songs: self.filenames.add(song[FIELD_SONGNAME]) def get_fingerprinted_songs(self) -> List[Dict[str, any]]: """ To pull all fingerprinted songs from the database. :return: a list of fingerprinted audios from the database. """ return self.db.get_songs() def delete_songs_by_id(self, song_ids: List[int]) -> None: """ Deletes all audios given their ids. :param song_ids: song ids to delete from the database. """ self.db.delete_songs_by_id(song_ids) def fingerprint_files(self, s3_keys: list, bucket: str, nprocesses: int = None) -> None: """ Fingerprint a list of S3Files :param s3_keys: list of s3 files to process. :param bucket: bucket that the files belong to. :param nprocesses: amount of processes to fingerprint the files within the directory. """ # Try to use the maximum amount of processes if not given. try: nprocesses = nprocesses or multiprocessing.cpu_count() except NotImplementedError: nprocesses = 1 else: nprocesses = 1 if nprocesses <= 0 else nprocesses print('Number of processes {}'.format(nprocesses)) pool = multiprocessing.Pool(nprocesses) # Prepare _fingerprint_worker input worker_input = list( zip( s3_keys, [bucket] * len(s3_keys), [self.limit] * len(s3_keys) ) ) # Send off our tasks iterator = pool.imap_unordered(Dejavu._fingerprint_worker, worker_input) # Loop till we have all of them while True: try: s3_key, hashes, file_hash = next(iterator) except multiprocessing.TimeoutError: continue except StopIteration: break except Exception: print('Failed fingerprinting {}') # Print traceback because we can't reraise it here traceback.print_exc(file=sys.stdout) else: song_name, extension = decoder.get_audio_name_from_path(s3_key) sid = self.db.insert_song(song_name, file_hash, len(hashes), extension) self.db.insert_hashes(sid, hashes) self.db.set_song_fingerprinted(sid) Dejavu.update_file_metadata(s3_key, bucket) pool.close() pool.join() return True def fingerprint_file(self, bucket: str, key: str) -> int: """ Given a path to a file the method generates hashes for it and stores them in the database for later be queried. :param bucket: bucket of the file. :param key: path to the file. :return: the song id. """ s3_object = s3.Object(bucket_name=bucket, key=key) s3_file = S3File(s3_object) song_name, extension = decoder.get_audio_name_from_path(key) hashes, file_hash = Dejavu.get_file_fingerprints(s3_file, self.limit) sid = self.db.insert_song(song_name, file_hash, len(hashes), extension) self.db.insert_hashes(sid, hashes) self.db.set_song_fingerprinted(sid) return sid @staticmethod def generate_fingerprints(samples: List[int], Fs=DEFAULT_FS) -> Tuple[List[Tuple[str, int]], float]: f""" Generate the fingerprints for the given sample data (channel). :param samples: list of ints which represents the channel info of the given audio file. :param Fs: sampling rate which defaults to {DEFAULT_FS}. :return: a list of tuples for hash and its corresponding offset, together with the generation time. """ t = time() hashes = fingerprint(samples, Fs=Fs) fingerprint_time = time() - t return hashes, fingerprint_time def find_matches(self, hashes: List[Tuple[str, int]], existing_song_id: int = None) -> Tuple[List[Tuple[int, int]], Dict[str, int], float]: """ Finds the corresponding matches on the fingerprinted audios for the given hashes. :param hashes: list of tuples for hashes and their corresponding offsets :return: a tuple containing the matches found against the db, a dictionary which counts the different hashes matched for each song (with the song id as key), and the time that the query took. """ t = time() matches, dedup_hashes = self.db.return_matches(hashes, existing_song_id) query_time = time() - t return matches, dedup_hashes, query_time def align_matches(self, matches: List[Tuple[int, int]], dedup_hashes: Dict[str, int], queried_hashes: int) -> List[Dict[str, any]]: """ Finds hash matches that align in time with other matches and finds consensus about which hashes are "true" signal from the audio. :param matches: matches from the database :param dedup_hashes: dictionary containing the hashes matched without duplicates for each song (key is the song id). :param queried_hashes: amount of hashes sent for matching against the db :param topn: number of results being returned back. :return: a list of dictionaries (based on topn) with match information. """ # count offset occurrences per song and keep only the maximum ones. sorted_matches = sorted(matches, key=lambda m: (m[0], m[1])) counts = [(*key, len(list(group))) for key, group in groupby(sorted_matches, key=lambda m: (m[0], m[1]))] songs_matches = sorted( [max(list(group), key=lambda g: g[2]) for key, group in groupby(counts, key=lambda count: count[0])], key=lambda count: count[2], reverse=True ) song_ids = [song_id for song_id, _, _ in songs_matches] songs_result = [] songs = self.db.get_songs_by_id(song_ids) for idx, song in enumerate(songs): song_id = song.get(SONG_ID, None) song_name = song.get(SONG_NAME, None) song_hashes = song.get(FIELD_TOTAL_HASHES, None) _, offset, matching_count = songs_matches[idx] nseconds = round(float(offset) / DEFAULT_FS * DEFAULT_WINDOW_SIZE * DEFAULT_OVERLAP_RATIO, 5) hashes_matched = dedup_hashes[song_id] song = { SONG_ID: song_id, SONG_NAME: song_name.encode('utf8'), INPUT_HASHES: queried_hashes, FINGERPRINTED_HASHES: song_hashes, HASHES_MATCHED: hashes_matched, # Percentage regarding hashes matched vs hashes from the input. INPUT_CONFIDENCE: round(hashes_matched / queried_hashes, 2), # Percentage regarding hashes matched vs hashes fingerprinted in the db. FINGERPRINTED_CONFIDENCE: round(hashes_matched / song_hashes, 2), OFFSET: offset, OFFSET_SECS: nseconds, FIELD_FILE_SHA1: song.get(FIELD_FILE_SHA1, None).encode('utf8') } songs_result.append(song) songs_result.sort(key=lambda x: x[FINGERPRINTED_CONFIDENCE], reverse=True) return songs_result def recognize(self, recognizer, *options, **kwoptions) -> Dict[str, any]: r = recognizer(self) return r.recognize(*options, **kwoptions) @staticmethod def _fingerprint_worker(arguments): # Pool.imap sends arguments as tuples so we have to unpack them ourself. s3_key, bucket, limit = arguments print('Fingerprinting file {}'.format(s3_key)) s3_object = s3.Object(bucket_name=bucket, key=s3_key) s3_file = S3File(s3_object) fingerprints, file_hash = Dejavu.get_file_fingerprints(s3_file, limit) return s3_key, fingerprints, file_hash @staticmethod def get_file_fingerprints(file: S3File, limit: int): channels, fs, file_hash = decoder.read(file, limit) key = file.s3_object.key fingerprints = set() channel_amount = len(channels) for channeln, channel in enumerate(channels, start=1): print('Fingerprinting channel {}/{} for {}'.format(channeln, channel_amount, key)) hashes = fingerprint(channel, Fs=fs) print('Finished channel {}/{} for {}'.format(channeln, channel_amount, key)) fingerprints |= set(hashes) print('Finished fingerprinting {}'.format(key)) return fingerprints, file_hash @staticmethod def update_file_metadata(key, bucket, state='1'): print("Updating metadata {}".format(key)) response = client.head_object( Bucket=bucket, Key=key ) metadata = response['Metadata'] metadata['fingerprinted'] = state copy_source = { 'Bucket': bucket, 'Key': key } s3.meta.client.copy( copy_source, bucket, key, ExtraArgs={ 'ContentType': response['ContentType'], 'Metadata': metadata, 'MetadataDirective': 'REPLACE' }, SourceClient=client ) print('Finished metadata state {} update {}'.format(state, key))