incremental waveform generation (kind-of) working

This commit is contained in:
Drew Weymouth
2025-07-26 15:33:22 -07:00
parent 88b7f62e50
commit c521c617ec
17 changed files with 365 additions and 160 deletions
+274 -71
View File
@@ -3,6 +3,7 @@ package backend
import (
"context"
"errors"
"fmt"
"image"
"image/color"
"io"
@@ -10,16 +11,18 @@ import (
"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 WaveformData struct {
Peak [1024]byte
RMS [1024]byte
type WaveformImageGenerator struct {
audioCache *AudioCache
}
type WaveformImage = image.NRGBA
@@ -28,27 +31,165 @@ func NewWaveformImage() *WaveformImage {
return image.NewNRGBA(image.Rect(0, 0, 1024, 32))
}
func GenerateWaveformImage(data *WaveformData, imgbuf *WaveformImage, c color.Color) {
centerY := imgbuf.Rect.Dy() / 2 // 16
type WaveformImageJob struct {
ItemID string
lock sync.Mutex
img *WaveformImage
err error
progress int // first invalid pixel in X direction
cancel func()
}
func (w *WaveformImageJob) Cancel() {
if w != nil && w.cancel != nil {
w.cancel()
}
}
func (w *WaveformImageJob) Done() bool {
w.lock.Lock()
defer w.lock.Unlock()
return w.err != nil || w.progress >= w.img.Bounds().Dx()
}
func (w *WaveformImageJob) Err() error {
w.lock.Lock()
defer w.lock.Unlock()
return w.err
}
func (w *WaveformImageJob) Get() *WaveformImage {
if w.Done() {
log.Println("returning image directly")
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 := 0; y < height; y++ {
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.PathForCachedOrDownloadingFile(job.ItemID)
// wait for file to begin downloading if not already
for path == "" {
time.Sleep(10 * time.Millisecond)
if e := ctx.Err(); e != nil {
job.setError(e)
return
}
path = w.audioCache.PathForCachedOrDownloadingFile(job.ItemID)
}
dir := filepath.Dir(path)
transcodeFile := filepath.Join(dir, filepath.Base(path)+"_waveform.wav")
fileDone := func() bool {
return w.audioCache.PathForCachedFile(job.ItemID) != ""
}
// If file isn't fully downloaded from server,
// stream it to MPV via HTTP 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()
log.Println("streaming file to MPV at ", path)
go srv.Serve()
}
// Start converting the file to WAV for analysis
var wavConvertDone bool
go func() {
err := convertToWav(ctx, 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(10 * 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 {
log.Println("error analyzing wav", err.Error())
job.setError(err)
}
}()
// Start generating the waveform image
go generateWaveformImage(ctx, data, job)
}()
return job
}
type waveformData struct {
Peak [1024]byte
RMS [1024]byte
progress int // first invalid index for Peak/RMS data
}
func generateWaveformImage(ctx context.Context, data *waveformData, job *WaveformImageJob) {
centerY := job.img.Rect.Dy() / 2 // 16
top := centerY - 1
bottom := centerY
// Convert the input color to RGBA
r, g, b, _ := c.RGBA()
opaqueColor := color.NRGBA{
R: uint8(r >> 8),
G: uint8(g >> 8),
B: uint8(b >> 8),
A: 255,
}
translucentColor := color.NRGBA{
R: uint8(r >> 8),
G: uint8(g >> 8),
B: uint8(b >> 8),
A: 128, // 50% opacity
}
opaqueColor := color.NRGBA{R: 255, G: 255, B: 255, A: 255}
translucentColor := color.NRGBA{R: 255, G: 255, B: 255, A: 128}
for x := 0; x < 1024; x++ {
if data.progress <= x {
time.Sleep(10 * time.Millisecond)
if ctx.Err() != nil {
return // expired
}
}
rms := float64(data.RMS[x]) / 255.0
peak := float64(data.Peak[x]) / 255.0
@@ -56,80 +197,110 @@ func GenerateWaveformImage(data *WaveformData, imgbuf *WaveformImage, c color.Co
peakPixels := int((peak - rms) * 16)
// Always draw at least 2 center pixels
setPixel(imgbuf, x, top, opaqueColor)
setPixel(imgbuf, x, bottom, opaqueColor)
setPixel(job.img, x, top, opaqueColor)
setPixel(job.img, x, bottom, opaqueColor)
// Draw RMS pixels (solid)
for i := 1; i <= rmsPixels; i++ {
setPixel(imgbuf, x, top-i, opaqueColor)
setPixel(imgbuf, x, bottom+i, opaqueColor)
setPixel(job.img, x, top-i, opaqueColor)
setPixel(job.img, x, bottom+i, opaqueColor)
}
// Draw Peak extension (translucent)
for i := 1; i <= peakPixels; i++ {
setPixel(imgbuf, x, top-rmsPixels-i, translucentColor)
setPixel(imgbuf, x, bottom+rmsPixels+i, translucentColor)
setPixel(job.img, x, top-rmsPixels-i, translucentColor)
setPixel(job.img, x, bottom+rmsPixels+i, translucentColor)
}
job.progress = x + 1
}
log.Println("done generating image")
}
func GetWaveformDataForFile(ctx context.Context, fpath string, fileIsDone func() bool) (*WaveformData, error) {
dir := filepath.Dir(fpath)
transcodeFile := filepath.Join(dir, filepath.Base(fpath)+"_waveform.wav")
if !fileIsDone() {
srv, err := util.NewFileStreamerServer(fpath, fileIsDone)
if err != nil {
return nil, err
}
fpath = srv.Addr()
log.Println("streaming file to MPV at ", fpath)
go srv.Serve()
func analyzeWavFile(ctx context.Context, transcodeFile string, data *waveformData, millisecs int64, fileDone func() bool) error {
if fileDone() {
log.Println("Analyzing completely written file!!")
}
err := convertToWav(ctx, fpath, transcodeFile)
if err != nil {
return nil, err
}
f, err := os.Open(transcodeFile)
if err != nil {
log.Println("error opening transcoded file")
return nil, err
return err
}
defer f.Close()
defer os.Remove(transcodeFile)
//defer os.Remove(transcodeFile)
decoder := wav.NewDecoder(f)
reader := trackingReader{rs: f}
decoder := wav.NewDecoder(&reader)
if !decoder.IsValidFile() {
return nil, errors.New("invalid wav file")
}
dur, err := decoder.Duration()
if err != nil {
return nil, err
return errors.New("invalid wav file")
}
format := decoder.Format()
totalSamples := format.SampleRate * int(dur.Milliseconds()) / 1000
totalSamples := format.SampleRate * int(millisecs) / 1000
samplesPerChunk := totalSamples / 1024
if err := decoder.FwdToPCM(); err != nil {
return nil, err
return err
}
buf := &audio.IntBuffer{Data: make([]int, 4096)}
data := &WaveformData{}
curChunk := 0
chunkSamples := make([]float32, 0, samplesPerChunk)
// file read loop
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
if !fileDone() {
// 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()
readableBytes := currentSize - reader.Pos() // how many bytes are still available
// Estimate how many samples we can read
bytesPerSample := int64(2 * format.NumChannels) // 16-bit = 2 bytes per channel
maxSamples := int(readableBytes / bytesPerSample)
if maxSamples <= 0 {
// Wait for more data to be written to file
time.Sleep(10 * time.Millisecond)
continue
}
// Resize buffer to fit only whats safe
safeSamples := maxSamples
if safeSamples > cap(buf.Data) {
safeSamples = 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 fileDone() {
data.progress = 1024 // set progress to done
}
if err == io.EOF && !fileDone() {
return errors.New("WAV read got premature EOF")
}
break
}
if err != nil {
return data, err
return err
}
// Process samples
@@ -150,6 +321,7 @@ func GetWaveformDataForFile(ctx context.Context, fpath string, fileIsDone func()
data.RMS[curChunk] = float32ToByte(rms)
}
curChunk++
data.progress = curChunk
chunkSamples = chunkSamples[:0]
if curChunk >= 1024 {
break
@@ -158,14 +330,24 @@ func GetWaveformDataForFile(ctx context.Context, fpath string, fileIsDone func()
}
}
// Optionally fill the last chunk if it's partially filled
// analyze the last chunk if it's partially filled with samples
if curChunk < 1024 && len(chunkSamples) > 0 {
peak, rms := computePeakAndRMS(chunkSamples)
data.Peak[curChunk] = float32ToByte(peak)
data.RMS[curChunk] = float32ToByte(rms)
data.progress = curChunk + 1
}
return data, nil
log.Println("final chunk is", curChunk)
return nil
}
func (j *WaveformImageJob) setError(err error) {
j.lock.Lock()
defer j.lock.Unlock()
j.err = err
}
func computePeakAndRMS(chunk []float32) (peak float32, rms float32) {
@@ -212,18 +394,7 @@ func convertToWav(ctx context.Context, inPath, outPath string) error {
m.Command([]string{"loadfile", inPath, "replace"})
return mpvWaitForIdle(ctx, m)
}
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
}
func mpvWaitForIdle(ctx context.Context, m *mpv.Mpv) error {
// Wait for MPV idle or ctx expiry
for {
select {
case <-ctx.Done():
@@ -235,10 +406,42 @@ func mpvWaitForIdle(ctx context.Context, m *mpv.Mpv) error {
}
// use small timeout to allow detecting ctx expiry
// without too much delay
e := m.WaitEvent(0.1 /*timeout seconds*/)
e := m.WaitEvent(0.05 /*timeout seconds*/)
if e.Event_Id == mpv.EVENT_IDLE {
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
}
// wrap an io.ReadSeeker with support for tracking bytes read (Pos())
type trackingReader struct {
rs io.ReadSeeker
pos int64
}
func (t *trackingReader) Read(p []byte) (int, error) {
n, err := t.rs.Read(p)
t.pos += int64(n)
return n, err
}
func (t *trackingReader) Seek(offset int64, whence int) (int64, error) {
newPos, err := t.rs.Seek(offset, whence)
if err == nil {
t.pos = newPos
}
return newPos, err
}
func (t *trackingReader) Pos() int64 {
return t.pos
}