Files
supersonic/backend/waveformimage.go
T
2025-09-17 22:30:49 +02:00

463 lines
11 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package backend
import (
"context"
"errors"
"fmt"
"image"
"image/color"
"io"
"log"
"math"
"os"
"path/filepath"
"sync"
"time"
"github.com/dweymouth/supersonic/backend/mediaprovider"
"github.com/dweymouth/supersonic/backend/util"
"github.com/go-audio/audio"
"github.com/go-audio/wav"
"github.com/supersonic-app/go-mpv"
)
type WaveformImageGenerator struct {
audioCache *AudioCache
}
type WaveformImage = image.NRGBA
func NewWaveformImage() *WaveformImage {
img := image.NewNRGBA(image.Rect(0, 0, 1024, 48))
centerTop := img.Rect.Dy() / 2 // 24
centerBottom := centerTop - 1 // 23
// color in center line
for x := 0; x < img.Bounds().Dx(); x++ {
setPixel(img, x, centerTop, color.NRGBA{R: 255, G: 255, B: 255, A: 255})
setPixel(img, x, centerBottom, color.NRGBA{R: 255, G: 255, B: 255, A: 255})
}
return img
}
type WaveformImageJob struct {
ItemID string
lock sync.Mutex
img *WaveformImage
err error
progress int // first invalid pixel in X direction
done bool
cancel func()
canceled bool
}
func (w *WaveformImageJob) Cancel() {
if w != nil {
w.canceled = true
if w.cancel != nil {
w.cancel()
}
}
}
func (w *WaveformImageJob) Canceled() bool {
return w.canceled
}
func (w *WaveformImageJob) Done() bool {
w.lock.Lock()
defer w.lock.Unlock()
return w.err != nil || w.done
}
func (w *WaveformImageJob) Err() error {
w.lock.Lock()
defer w.lock.Unlock()
return w.err
}
func (w *WaveformImageJob) Get() *WaveformImage {
if w.Done() {
return w.img
}
// return a new *WaveformImage with data copied
// from the valid region of w.img
height := w.img.Bounds().Dy()
result := NewWaveformImage()
// Copy each scanline from w.img to result
for y := range height {
srcOffset := w.img.PixOffset(0, y)
dstOffset := result.PixOffset(0, y)
copy(result.Pix[dstOffset:dstOffset+w.progress*4], w.img.Pix[srcOffset:srcOffset+w.progress*4])
}
return result
}
func NewWaveformImageGenerator(cache *AudioCache) *WaveformImageGenerator {
return &WaveformImageGenerator{audioCache: cache}
}
func (w *WaveformImageGenerator) StartWaveformGeneration(item *mediaprovider.Track) *WaveformImageJob {
ctx, cancel := context.WithCancel(w.audioCache.rootCtx)
job := &WaveformImageJob{
img: NewWaveformImage(),
ItemID: item.ID,
cancel: cancel,
}
// Set up a pipeline of concurrent tasks that need to complete to generate
// a waveform image:
// 1. Begin downloading the file from the server
// 2. Begin transcoding it to WAV
// 3. Begin analyzing the resulting WAV file
// 4. Begin generating the image from the analysis data
go func() {
path := w.audioCache.ObtainReferenceToFile(job.ItemID)
// wait for file to begin downloading if not already
for path == "" {
time.Sleep(50 * time.Millisecond)
if e := ctx.Err(); e != nil {
job.setError(e)
return
}
path = w.audioCache.ObtainReferenceToFile(job.ItemID)
}
// and wait for content to begin being written
for {
if s, err := os.Stat(path); err == nil && s.Size() > 0 {
break
}
time.Sleep(50 * time.Millisecond)
if e := ctx.Err(); e != nil {
job.setError(e)
return
}
}
dir := filepath.Dir(path)
var transcodeFile string
for i := 0; true; i++ {
if i > 0 {
transcodeFile = filepath.Join(dir, fmt.Sprintf("%s_waveform_%d.wav", filepath.Base(path), i))
} else {
transcodeFile = filepath.Join(dir, filepath.Base(path)+"_waveform.wav")
}
if _, err := os.Stat(transcodeFile); os.IsNotExist(err) {
break // found a suitable filename that doesn't exist
}
}
fileDone := func() bool {
return w.audioCache.IsFullyDownloaded(job.ItemID)
}
// If file isn't fully downloaded from server,
// stream it to MPV via fifo so it doesn't possibly
// terminate the conversion to WAV early encountering EOF
if !fileDone() {
srv, err := util.NewFileStreamerServer(path, fileDone)
if err != nil {
job.setError(err)
return
}
path = srv.Addr()
go srv.Serve()
time.Sleep(10 * time.Millisecond) // make sure server has time to come up
}
// Start converting the file to WAV for analysis
var wavConvertDone bool
go func() {
err := w.convertToWav(ctx, job.ItemID, path, transcodeFile)
wavConvertDone = true
if err != nil {
job.setError(err)
}
}()
// Wait for transcoded WAV file to begin being written
for {
if s, err := os.Stat(transcodeFile); err == nil && s.Size() > 0 {
break
}
time.Sleep(50 * time.Millisecond)
if e := ctx.Err(); e != nil {
job.setError(e)
return
}
}
// Start analyzing the converted wav file
data := &waveformData{}
go func() {
err := analyzeWavFile(ctx, transcodeFile, data, item.Duration.Milliseconds(), func() bool { return wavConvertDone })
if err != nil {
job.setError(err)
}
data.done = true
}()
// Start generating the waveform image
go func() {
generateWaveformImage(ctx, data, job)
job.done = true
}()
}()
return job
}
type waveformData struct {
Peak [1024]byte
RMS [1024]byte
progress int // first invalid index for Peak/RMS data
done bool
}
func generateWaveformImage(ctx context.Context, data *waveformData, job *WaveformImageJob) {
centerY := job.img.Rect.Dy() / 2 // 24
top := centerY - 1 // 23
bottom := centerY // 24
opaqueColor := color.NRGBA{R: 255, G: 255, B: 255, A: 255}
translucentColor := color.NRGBA{R: 255, G: 255, B: 255, A: 128}
for x := range 1024 {
for data.progress <= x {
if data.done {
return
}
if ctx.Err() != nil {
return // expired
}
time.Sleep(50 * time.Millisecond)
}
rmsPixels := int(data.RMS[x]) * centerY / 255
peakPixels := int(data.Peak[x]) * centerY / 255
// Always draw at least 2 center pixels
setPixel(job.img, x, top, opaqueColor)
setPixel(job.img, x, bottom, opaqueColor)
// Draw RMS pixels (solid)
for i := 1; i < rmsPixels; i++ {
setPixel(job.img, x, top-i, opaqueColor)
setPixel(job.img, x, bottom+i, opaqueColor)
}
// Draw Peak extension (translucent)
for i := max(1, rmsPixels); i < peakPixels; i++ {
setPixel(job.img, x, top-i, translucentColor)
setPixel(job.img, x, bottom+i, translucentColor)
}
job.progress = x + 1
}
}
// assumes mono, 16 bit
func analyzeWavFile(ctx context.Context, transcodeFile string, data *waveformData, millisecs int64, fileDone func() bool) error {
f, err := os.Open(transcodeFile)
if err != nil {
return err
}
defer f.Close()
defer os.Remove(transcodeFile)
decoder := wav.NewDecoder(f)
if !decoder.IsValidFile() {
return errors.New("invalid wav file")
}
format := decoder.Format()
totalSamples := format.SampleRate * int(millisecs) / 1000
samplesPerChunk := totalSamples / 1024
if err := decoder.FwdToPCM(); err != nil {
return err
}
buf := &audio.IntBuffer{Data: make([]int, 4096)}
curChunk := 0
chunkSamples := make([]float64, 0, samplesPerChunk)
bytesPerSample := int64(2 * format.NumChannels) // 16-bit = 2 bytes per channel
// file read loop
doneReading := false
for !doneReading {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
fileIsDone := fileDone()
if !fileIsDone {
// Check how many samples we can safely read without encountering EOF
// and adjust read buffer size accordingly
stat, err := f.Stat()
if err != nil {
return fmt.Errorf("stat failed: %w", err)
}
currentSize := stat.Size()
// how many bytes can we read without nearing EOF
readableBytes := currentSize - int64(samplesPerChunk)*int64(curChunk)*bytesPerSample - 8192 // buffer for safety
// Estimate how many samples we can read
maxSamples := int(readableBytes / bytesPerSample)
if maxSamples <= 0 {
// Wait for more data to be written to file
time.Sleep(50 * time.Millisecond)
continue
}
// Resize buffer to fit only whats safe
safeSamples := min(maxSamples, cap(buf.Data))
buf.Data = buf.Data[:safeSamples]
} else {
// File is done being written, resize read buf to the max
buf.Data = buf.Data[:cap(buf.Data)]
}
n, err := decoder.PCMBuffer(buf)
if n == 0 || err == io.EOF {
if fileIsDone {
doneReading = true
}
if err == io.EOF && !fileDone() {
return errors.New("WAV read got premature EOF")
}
}
if err != nil {
return err
}
// Process samples
for i := range n {
sample := float64(buf.Data[i]) / float64(1<<15) // Normalize to [-1, 1]
chunkSamples = append(chunkSamples, sample)
if len(chunkSamples) >= samplesPerChunk {
if curChunk < 1024 {
peak, rms := computePeakAndRMS(chunkSamples)
data.Peak[curChunk] = float64ToByte(peak)
data.RMS[curChunk] = float64ToByte(rms)
}
curChunk++
data.progress = curChunk
chunkSamples = chunkSamples[:0]
if curChunk >= 1024 {
break
}
}
}
}
// analyze the last chunk if it's partially filled with samples
if curChunk < 1024 && len(chunkSamples) > 0 {
peak, rms := computePeakAndRMS(chunkSamples)
data.Peak[curChunk] = float64ToByte(peak)
data.RMS[curChunk] = float64ToByte(rms)
data.progress = curChunk + 1
}
return nil
}
func (j *WaveformImageJob) setError(err error) {
j.lock.Lock()
defer j.lock.Unlock()
j.err = err
}
func computePeakAndRMS(chunk []float64) (peak float64, rms float64) {
var sumSquares float64
peak = 0.0
for _, v := range chunk {
if v > peak {
peak = v
} else if v < -peak {
peak = -v
}
sumSquares += float64(v * v)
}
rms = math.Sqrt(sumSquares / float64(len(chunk)))
return peak, rms
}
func float64ToByte(val float64) byte {
if val > 1.0 {
val = 1.0
}
if val < 0.0 {
val = 0.0
}
return byte(val * 255)
}
func (w *WaveformImageGenerator) convertToWav(ctx context.Context, id, inPath, outPath string) error {
m := mpv.Create()
m.SetOptionString("video", "no")
m.SetOptionString("audio-display", "no")
m.SetOptionString("terminal", "no")
m.SetOptionString("idle", "yes")
m.SetOptionString("ao-pcm-file", outPath)
m.SetOptionString("ao", "pcm")
m.SetOption("volume", mpv.FORMAT_INT64, 100)
// no need to preserve full sample resolution just for waveform image
// let's make less data to process and smaller on-disk file
m.SetOption("audio-samplerate", mpv.FORMAT_INT64, 22050)
m.SetOptionString("audio-channels", "mono")
m.SetOptionString("audio-format", "s16")
if err := m.Initialize(); err != nil {
return err
}
defer m.TerminateDestroy()
m.Command([]string{"loadfile", inPath, "replace"})
defer w.audioCache.ReleaseReferenceToFile(id)
// Wait for MPV idle or ctx expiry
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
// use small timeout to allow detecting ctx expiry
// without too much delay
e := m.WaitEvent(0.05 /*timeout seconds*/)
if e.Event_Id == mpv.EVENT_IDLE {
if _, err := os.Stat(outPath); os.IsNotExist(err) {
log.Printf("WARNING! file %s does not exist after MPV convert", outPath)
}
return nil
}
ia := m.GetPropertyString("idle-active")
if ia == "yes" || ia == "true" {
if _, err := os.Stat(outPath); os.IsNotExist(err) {
log.Printf("WARNING! file %s does not exist after MPV convert", outPath)
}
return nil
}
}
}
}
func setPixel(img *image.NRGBA, x, y int, c color.NRGBA) {
offset := img.PixOffset(x, y)
img.Pix[offset+0] = c.R
img.Pix[offset+1] = c.G
img.Pix[offset+2] = c.B
img.Pix[offset+3] = c.A
}