Implement reassignment method for sharp time-frequency localization

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>

src/spectrogram.c

@@ -34,6 +34,9 @@
 #define MAX_SAMPLE_RATE 48000
 #define LOUDNESS_FLOOR_DB -80.0f
 
+// Reassignment method for sharper time-frequency localization
+#define USE_REASSIGNMENT 1
+
 // Colormap types
 typedef enum {
     COLORMAP_GRAYS = 0,
@@ -359,7 +362,7 @@ static void FreeSignal(AudioSignal* signal)
 }
 
 // ============================================================================
-// Spectrogram Generation
+// Spectrogram Generation with Reassignment
 // ============================================================================
 
 static void GenerateSpectrogramTexture(StftResult* stft, Image* image, Texture2D* texture)
@@ -370,21 +373,107 @@ static void GenerateSpectrogramTexture(StftResult* stft, Image* image, Texture2D
     *image = GenImageColor(width, height, BLACK);
     Color* pixels = (Color*)image->data;
     
+    // Initialize to black
+    for (int i = 0; i < width * height; i++) {
+        pixels[i] = BLACK;
+    }
+
+    // Find max amplitude for normalization
     float maxAmplitude = 0.0001f;
     for (int seg = 0; seg < stft->numSegments; seg++)
         for (int bin = 0; bin < stft->segments[seg].numBins; bin++)
-            if (stft->segments[seg].spectrum[bin].amplitude > maxAmplitude) maxAmplitude = stft->segments[seg].spectrum[bin].amplitude;
+            if (stft->segments[seg].spectrum[bin].amplitude > maxAmplitude) 
+                maxAmplitude = stft->segments[seg].spectrum[bin].amplitude;
+
+    float hopSize = (float)(stft->segments[0].sampleOffset > 0 ? 
+                   stft->segments[1].sampleOffset - stft->segments[0].sampleOffset : 
+                   stft->segments[0].sampleCount);
+    float freqPerBin = (float)stft->sampleRate / (height * 2 - 2);
+    
+    // Create a floating-point accumulation buffer for reassignment
+    float* accumBuffer = (float*)calloc(width * height, sizeof(float));
     
     for (int seg = 0; seg < width; seg++) {
+        float segmentTime = (seg * hopSize) / (float)stft->sampleRate;
+        
         for (int bin = 0; bin < height; bin++) {
             float amplitude = stft->segments[seg].spectrum[bin].amplitude;
-            float db = AmplitudeToDecibels(amplitude);
+            float phase = stft->segments[seg].spectrum[bin].phase;
+            
+            if (amplitude < 0.0001f) continue;
+            
+            float reassignedFreq = bin * freqPerBin;
+            float reassignedTime = segmentTime;
+            
+#if USE_REASSIGNMENT
+            // ===== Reassignment Method =====
+            // Estimate instantaneous frequency from phase derivative over time
+            if (seg > 0 && seg < width - 1) {
+                float prevPhase = stft->segments[seg-1].spectrum[bin].phase;
+                float nextPhase = stft->segments[seg+1].spectrum[bin].phase;
+                
+                // Phase difference (unwrapped)
+                float phaseDiff = nextPhase - prevPhase;
+                // Unwrap to [-pi, pi]
+                while (phaseDiff > M_PI) phaseDiff -= 2.0f * M_PI;
+                while (phaseDiff < -M_PI) phaseDiff += 2.0f * M_PI;
+                
+                // Instantaneous frequency deviation
+                float expectedPhaseShift = 2.0f * M_PI * bin * hopSize / (height * 2 - 2);
+                float phaseDeviation = phaseDiff - expectedPhaseShift;
+                float instantFreqDev = phaseDeviation * stft->sampleRate / (2.0f * M_PI * hopSize);
+                
+                reassignedFreq = bin * freqPerBin + instantFreqDev;
+            }
+            
+            // Estimate group delay from phase derivative over frequency
+            if (bin > 0 && bin < height - 1) {
+                float prevPhaseAdj = stft->segments[seg].spectrum[bin-1].phase;
+                float nextPhaseAdj = stft->segments[seg].spectrum[bin+1].phase;
+                
+                float phaseGrad = nextPhaseAdj - prevPhaseAdj;
+                // Unwrap
+                while (phaseGrad > M_PI) phaseGrad -= 2.0f * M_PI;
+                while (phaseGrad < -M_PI) phaseGrad += 2.0f * M_PI;
+                
+                // Group delay (time correction)
+                float groupDelay = -phaseGrad / (2.0f * M_PI * freqPerBin);
+                reassignedTime = segmentTime + groupDelay / (float)stft->sampleRate;
+            }
+            
+            // Clamp to valid range
+            if (reassignedFreq < 0) reassignedFreq = 0;
+            if (reassignedFreq >= stft->sampleRate / 2.0f) reassignedFreq = stft->sampleRate / 2.0f - 1;
+#endif
+            
+            // Map reassigned coordinates to pixel indices
+            int reassignedBin = (int)(reassignedFreq / freqPerBin);
+            int reassignedSeg = (int)((reassignedTime * stft->sampleRate) / hopSize);
+            
+            // Clamp to texture bounds
+            if (reassignedBin < 0) reassignedBin = 0;
+            if (reassignedBin >= height) reassignedBin = height - 1;
+            if (reassignedSeg < 0) reassignedSeg = 0;
+            if (reassignedSeg >= width) reassignedSeg = width - 1;
+            
+            // Accumulate amplitude at reassigned location
+            int pixelIndex = (height - 1 - reassignedBin) * width + reassignedSeg;
+            accumBuffer[pixelIndex] += amplitude;
+        }
+    }
+    
+    // Convert accumulation buffer to colors
+    for (int i = 0; i < width * height; i++) {
+        if (accumBuffer[i] > 0.0001f) {
+            float db = AmplitudeToDecibels(accumBuffer[i]);
             float normalized = (db - app.amplitudeFloorDb) / (app.amplitudeCeilingDb - app.amplitudeFloorDb);
             normalized = Clamp(normalized, 0.0f, 1.0f);
-            int pixelIndex = (height - 1 - bin) * width + seg;
-            pixels[pixelIndex] = GetColormapColor(normalized, app.colormap);
+            pixels[i] = GetColormapColor(normalized, app.colormap);
         }
     }
+    
+    free(accumBuffer);
+    
     if (texture->id != 0) UnloadTexture(*texture);
     *texture = LoadTextureFromImage(*image);
     SetTextureFilter(*texture, TEXTURE_FILTER_BILINEAR);
@@ -805,6 +894,9 @@ static void DrawInfo(Rectangle bounds)
         DrawText(TextFormat("Duration: %.2f sec", app.signal.duration), 10, y, fontSize, LIGHTGRAY); y += 20;
         DrawText(TextFormat("View: %.1f%%-%.1f%% (%.2f sec)", app.viewStart*100, app.viewEnd*100, (app.viewEnd-app.viewStart)*app.signal.duration), 10, y, fontSize, LIGHTGRAY); y += 20;
         DrawText(TextFormat("FFT: %d (%.1f Hz/bin, 75%% overlap)", app.fftSize, (float)app.signal.sampleRate / app.fftSize), 10, y, fontSize, LIGHTGRAY); y += 20;
+#if USE_REASSIGNMENT
+        DrawText("Reassignment: ON (sharp)", 10, y, fontSize, (Color){ 80, 255, 80, 255 }); y += 20;
+#endif
         DrawText(TextFormat("Max Freq: %.1f kHz", (float)app.signal.sampleRate / 2000.0f), 10, y, fontSize, LIGHTGRAY); y += 20;
         y += 10;