File size: 6,747 Bytes
e6ed91e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
"""

Simplified postprocessing for piano transcription

Robust and memory-efficient version for HuggingFace Spaces

"""

import numpy as np
import librosa
from scipy.signal import find_peaks
import logging

logger = logging.getLogger(__name__)

class MusicTranscriptionPostprocessor:
    """Simplified postprocessor for robust operation"""
    
    def __init__(self,

                 onset_threshold=0.3,

                 frame_threshold=0.3,

                 min_note_duration=0.05,

                 max_note_duration=8.0,

                 time_resolution=0.032):
        
        self.onset_threshold = onset_threshold
        self.frame_threshold = frame_threshold
        self.min_note_duration = min_note_duration
        self.max_note_duration = max_note_duration
        self.time_resolution = time_resolution

    def process_predictions(self, predictions):
        """Main processing function - simplified and robust"""
        try:
            logger.info("🎼 Processing model predictions...")
            
            # Handle different output formats
            if len(predictions) >= 3:
                onset_preds = predictions[0][0]
                frame_preds = predictions[1][0] 
                velocity_preds = predictions[2][0] if len(predictions) > 2 else None
            else:
                raise ValueError(f"Expected at least 3 model outputs, got {len(predictions)}")
            
            logger.info(f"📊 Prediction shapes: onset{onset_preds.shape}, frame{frame_preds.shape}")
            
            # Extract notes using simple but reliable method
            notes = self._extract_notes_simple(onset_preds, frame_preds, velocity_preds)
            
            # Clean up the notes
            cleaned_notes = self._clean_notes(notes)
            
            logger.info(f"✅ Extracted {len(cleaned_notes)} notes")
            return cleaned_notes
            
        except Exception as e:
            logger.error(f"❌ Postprocessing failed: {e}")
            return self._fallback_notes()
    
    def _extract_notes_simple(self, onset_preds, frame_preds, velocity_preds):
        """Simple but robust note extraction"""
        notes = []
        
        try:
            # Process each pitch
            for pitch_idx in range(min(88, onset_preds.shape[1])):
                onset_curve = onset_preds[:, pitch_idx]
                frame_curve = frame_preds[:, pitch_idx]
                
                # Find onset peaks
                peaks, _ = find_peaks(
                    onset_curve,
                    height=self.onset_threshold,
                    distance=max(1, int(0.05 / self.time_resolution))  # Min 50ms apart
                )
                
                # Create notes from peaks
                for peak in peaks:
                    # Find note duration using frame predictions
                    duration = self._find_note_duration(peak, frame_curve)
                    
                    if duration >= self.min_note_duration:
                        # Get velocity
                        velocity = self._get_velocity(peak, pitch_idx, velocity_preds)
                        
                        # Create note
                        midi_pitch = pitch_idx + 21  # Piano range starts at A0 (21)
                        note = {
                            "note_name": self._pitch_to_note_name(midi_pitch),
                            "time": float(peak * self.time_resolution),
                            "duration": float(duration),
                            "velocity": float(velocity),
                            "velocity_midi": int(min(127, max(1, velocity * 127))),
                            "pitch": int(midi_pitch),
                            "frequency": librosa.midi_to_hz(midi_pitch)
                        }
                        notes.append(note)
        
        except Exception as e:
            logger.error(f"Note extraction error: {e}")
        
        return notes
    
    def _find_note_duration(self, onset_frame, frame_curve):
        """Find note duration using frame predictions"""
        try:
            # Look for where the frame prediction drops below threshold
            remaining_frames = frame_curve[onset_frame:]
            
            # Find first point below threshold
            below_threshold = np.where(remaining_frames < self.frame_threshold)[0]
            
            if len(below_threshold) > 0:
                duration_frames = below_threshold[0]
            else:
                # Default duration if no clear ending
                duration_frames = min(int(0.5 / self.time_resolution), len(remaining_frames))
            
            duration = duration_frames * self.time_resolution
            return min(self.max_note_duration, max(self.min_note_duration, duration))
            
        except Exception:
            return 0.5  # Default duration
    
    def _get_velocity(self, onset_frame, pitch_idx, velocity_preds):
        """Get velocity for the note"""
        try:
            if velocity_preds is not None:
                raw_velocity = velocity_preds[onset_frame, pitch_idx]
                return float(np.clip(raw_velocity, 0.0, 1.0))
            else:
                return 0.8  # Default velocity
        except Exception:
            return 0.8
    
    def _clean_notes(self, notes):
        """Clean and filter notes"""
        if not notes:
            return notes
        
        # Sort by time
        notes.sort(key=lambda x: x["time"])
        
        # Remove duplicates and very short notes
        cleaned = []
        for note in notes:
            if (note["duration"] >= self.min_note_duration and 
                note["time"] >= 0 and
                21 <= note["pitch"] <= 108):  # Valid piano range
                cleaned.append(note)
        
        return cleaned
    
    def _pitch_to_note_name(self, pitch):
        """Convert MIDI pitch to note name"""
        note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        octave = (pitch // 12) - 1
        note = note_names[pitch % 12]
        return f"{note}{octave}"
    
    def _fallback_notes(self):
        """Fallback notes if processing fails"""
        logger.warning("Using fallback notes")
        return [
            {
                "note_name": "C4",
                "time": 0.0,
                "duration": 1.0,
                "velocity": 0.8,
                "velocity_midi": 80,
                "pitch": 60,
                "frequency": 261.63
            }
        ]