ahmed/acousticbrainz-ng
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Save data as image

Browse Source
This commit is contained in:
Ahmed Al-Taiar
2025-08-30 16:56:00 -04:00
parent 99c8d5aa41
commit 9e59b9a673
2 changed files with 182 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
__init__.py
View File

@@ -2,10 +2,12 @@ import os
import json
import subprocess
import hashlib
import zlib
import struct
import threading
import concurrent.futures
import tempfile
import re
from functools import partial
from typing import List, Tuple, Dict, Optional
@@ -19,6 +21,10 @@ from picard.track import Track
from picard.album import Album
from picard.ui.options import OptionsPage, register_options_page
from picard.util import thread
from picard.coverart.image import (
CoverArtImage,
CoverArtImageError,
)
from PyQt5 import QtWidgets, QtCore
_analysis_semaphore = None
@@ -380,6 +386,141 @@ class AcousticBrainzNG:
return True
def save_fingerprint(self, metadata: Dict, file_path: str, file_obj) -> bool:
if not self._check_optional_models():
log.error("Optional models not found")
return False
models_path = config.setting["acousticbrainz_ng_models_path"]
if not models_path:
log.error("Models path not configured")
return False
try:
output_path = self._generate_cache_folder(metadata, file_path)
if not output_path:
log.error("Failed to generate cache folder path")
return False
except Exception as e:
log.error(f"Error generating cache folder: {e}")
return False
try:
output_path = self._generate_cache_folder(metadata, file_path)
if not output_path:
log.error("Failed to generate cache folder path")
return False
except Exception as e:
log.error(f"Error generating cache folder: {e}")
return False
fingerprint_data = []
for key, value in metadata.items():
if key.lower().startswith("ab:hi:"):
try:
float_value = float(value)
if 0 <= float_value <= 1:
fingerprint_data.append(float_value)
except (ValueError, TypeError):
continue
if not fingerprint_data:
log.error("No valid fingerprint data found in metadata")
return False
if len(fingerprint_data) != 95:
log.error(f"Fingerprint expected exactly 95 values, got {len(fingerprint_data)}")
return False
fingerprint_file = os.path.join(output_path, "fingerprint.png")
try:
try:
import numpy as _np
except Exception as e:
log.error(f"numpy is required to generate fingerprint PNG: {e}")
return False
def _checksum_floats(values, n=5):
arr = _np.clip(_np.asarray(values, dtype=float).flatten(), 0.0, 1.0)
b = (arr * 65535).astype(_np.uint16).tobytes()
buf = hashlib.sha256(b).digest()
while len(buf) < n * 4:
buf += hashlib.sha256(buf).digest()
out = []
for i in range(n):
start = i * 4
u = struct.unpack(">I", buf[start:start+4])[0]
out.append(u / 0xFFFFFFFF)
return _np.array(out, dtype=float)
def _to_grayscale_uint8(arr):
a = _np.clip(_np.asarray(arr, dtype=float), 0.0, 1.0)
return (255 - _np.round(a * 255)).astype(_np.uint8)
def _png_write_grayscale(path, img8):
if img8.ndim != 2 or img8.dtype != _np.uint8:
raise ValueError("img8 must be a 2D numpy array of dtype uint8")
height, width = int(img8.shape[0]), int(img8.shape[1])
def _chunk(c_type, data):
chunk = struct.pack(">I", len(data)) + c_type + data
crc = zlib.crc32(c_type + data) & 0xFFFFFFFF
return chunk + struct.pack(">I", crc)
png_sig = b'\x89PNG\r\n\x1a\n'
ihdr = struct.pack(">IIBBBBB",
width, height,
8, # bit depth
0, # color type = 0 (grayscale)
0, # compression
0, # filter
0) # interlace
raw = bytearray()
for y in range(height):
raw.append(0)
raw.extend(img8[y].tobytes())
comp = zlib.compress(bytes(raw), level=9)
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as f:
f.write(png_sig)
f.write(_chunk(b'IHDR', ihdr))
f.write(_chunk(b'IDAT', comp))
f.write(_chunk(b'IEND', b''))
v = _np.clip(_np.asarray(fingerprint_data, dtype=float).flatten(), 0.0, 1.0)
base = _np.zeros(100, dtype=float)
base[:95] = v
base[95:] = _checksum_floats(v, n=5)
base = base.reshape((10, 10))
img8 = _to_grayscale_uint8(base)
_png_write_grayscale(fingerprint_file, img8)
fingerprint_url = f"file://{fingerprint_file.replace(os.sep, '/')}"
with open(fingerprint_file, "rb") as f:
fingerprint_data_bytes = f.read()
cover_art_image = CoverArtImage(url=fingerprint_url, data=fingerprint_data_bytes, comment=f"{PLUGIN_NAME} fingerprint", types=['other'], support_types=True)
file_obj.metadata.images.append(cover_art_image)
file_obj.metadata_images_changed.emit()
except Exception as e:
log.error(f"Failed to create fingerprint PNG: {e}")
return False
return True
@staticmethod
def _format_class(class_name: str) -> str:
return class_name.replace("non", "not").replace("_", " ").capitalize()
@@ -891,6 +1032,15 @@ def analyze_track(track: Track, album: Optional[Album] = None) -> Dict:
log.error(error_msg)
results['errors'].append(error_msg)
if config.setting["acousticbrainz_ng_save_fingerprint"]:
sf_result = acousticbrainz_ng.save_fingerprint(file.metadata, file.filename, file)
if not sf_result:
error_msg = f"Failed to save fingerprint for {file.filename}"
log.error(error_msg)
results['errors'].append(error_msg)
else:
file.metadata_images_changed.emit()
results['files_processed'] += 1
except Exception as e:
@@ -1037,7 +1187,7 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
options_layout = QtWidgets.QVBoxLayout(options_group)
self.analyze_optional_checkbox = QtWidgets.QCheckBox("Analyze optional MusicNN models", self)
self.analyze_optional_checkbox.setToolTip("Include optional MusicNN models in the analysis, currently unused unless raw values is enabled")
self.analyze_optional_checkbox.setToolTip("Include optional MusicNN models in the analysis")
self.save_raw_checkbox = QtWidgets.QCheckBox("Save raw values", self)
self.save_raw_checkbox.setToolTip("Save raw MusicNN numbers in the metadata")
@@ -1045,6 +1195,11 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
self.calculate_replaygain_checkbox = QtWidgets.QCheckBox("Calculate ReplayGain", self)
self.calculate_replaygain_checkbox.setToolTip("Calculate ReplayGain values for the track and album")
self.save_fingerprint_checkbox = QtWidgets.QCheckBox("Save fingerprint image", self)
self.save_fingerprint_checkbox.setToolTip("Save MusicNN data as an image, requires optional MusicNN values")
self.analyze_optional_checkbox.toggled.connect(self._update_fingerprint_state)
musicnn_workers_layout = QtWidgets.QHBoxLayout()
concurrent_analyses_layout = QtWidgets.QHBoxLayout()
rg_reference_layout = QtWidgets.QHBoxLayout()
@@ -1086,6 +1241,7 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
options_layout.addWidget(self.analyze_optional_checkbox)
options_layout.addWidget(self.save_raw_checkbox)
options_layout.addWidget(self.save_fingerprint_checkbox)
options_layout.addWidget(self.calculate_replaygain_checkbox)
options_layout.addLayout(concurrent_analyses_layout)
options_layout.addLayout(musicnn_workers_layout)
@@ -1169,6 +1325,11 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
layout.addStretch()
def _update_fingerprint_state(self, checked):
if not checked:
self.save_fingerprint_checkbox.setChecked(False)
self.save_fingerprint_checkbox.setEnabled(checked)
def _check_binaries(self, show_success=False) -> bool:
binaries_path = self.binaries_path_input.text()
if not binaries_path or not os.path.exists(binaries_path):
@@ -1277,6 +1438,16 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
self.rg_reference_input.setEnabled(self.calculate_replaygain_checkbox.isChecked())
fingerprint_setting = config.setting["acousticbrainz_ng_save_fingerprint"]
optional_setting = config.setting["acousticbrainz_ng_analyze_optional"] or False
if fingerprint_setting is None:
self.save_fingerprint_checkbox.setChecked(True if optional_setting else False)
else:
self.save_fingerprint_checkbox.setChecked(fingerprint_setting if optional_setting else False)
self._update_fingerprint_state(optional_setting)
self.musicnn_workers_input.setValue(config.setting["acousticbrainz_ng_max_musicnn_workers"] or 4)
self.concurrent_analyses_input.setValue(config.setting["acousticbrainz_ng_max_concurrent_analyses"] or 2)
self.rg_reference_input.setValue(config.setting["acousticbrainz_ng_replaygain_reference_loudness"] or -18)
@@ -1294,6 +1465,11 @@ class AcousticBrainzNGOptionsPage(OptionsPage):
config.setting["acousticbrainz_ng_save_raw"] = self.save_raw_checkbox.isChecked()
config.setting["acousticbrainz_ng_calculate_replaygain"] = self.calculate_replaygain_checkbox.isChecked()
if self.analyze_optional_checkbox.isChecked():
config.setting["acousticbrainz_ng_save_fingerprint"] = self.save_fingerprint_checkbox.isChecked()
else:
config.setting["acousticbrainz_ng_save_fingerprint"] = False
max_workers = max(1, min(self.musicnn_workers_input.value(), max(len(REQUIRED_MODELS), len(OPTIONAL_MODELS))))
config.setting["acousticbrainz_ng_max_musicnn_workers"] = max_workers

5
constants.py
View File

@@ -18,7 +18,7 @@ External dependencies:
</ul>
<strong>This plugin is CPU heavy!</strong>
"""
PLUGIN_VERSION = "1.0.0"
PLUGIN_VERSION = "1.1.0"
PLUGIN_API_VERSIONS = ["2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.13"]
PLUGIN_LICENSE = "GPL-2.0-or-later"
PLUGIN_LICENSE_URL = "https://www.gnu.org/licenses/gpl-2.0.html"
@@ -64,7 +64,8 @@ CONFIG_OPTIONS = [
IntOption("setting", "acousticbrainz_ng_replaygain_reference_loudness", -18),
BoolOption("setting", "acousticbrainz_ng_analyze_optional", False),
BoolOption("setting", "acousticbrainz_ng_save_raw", False),
BoolOption("setting", "acousticbrainz_ng_calculate_replaygain", True)
BoolOption("setting", "acousticbrainz_ng_calculate_replaygain", True),
BoolOption("setting", "acousticbrainz_ng_save_fingerprint", True)
]
GAIA_KEY_ALGORITHMS = ["edma", "krumhansl", "temperley"]