Update custom_interface_app.py

custom_interface_app.py

@@ -231,48 +231,65 @@ class ASR(Pretrained):
             return outputs
 
 
-    def classify_file_whisper_mkd_streaming(self, waveform, device):
-        # Load the audio file
-        # waveform, sr = librosa.load(path, sr=16000)
-
+    def classify_file_w2v2(self, waveform, device):
         # Get audio length in seconds
-        audio_length = len(waveform) / 16000
+        sr = 16000
+        audio_length = len(waveform) / sr
         
         if audio_length >= 20:
-            # split audio every 20 seconds
+            print(f"Audio is too long ({audio_length:.2f} seconds), splitting into segments")
+            # Detect non-silent segments
+
+            non_silent_intervals = librosa.effects.split(waveform, top_db=20)  # Adjust top_db for sensitivity
+
             segments = []
-            max_duration = 20 * 16000  # Maximum segment duration in samples (20 seconds)
-            num_segments = int(np.ceil(len(waveform) / max_duration))
-            start = 0
-            for i in range(num_segments):
-                end = start + max_duration
-                if end > len(waveform):
-                    end = len(waveform)
+            current_segment = []
+            current_length = 0
+            max_duration = 20 * sr  # Maximum segment duration in samples (20 seconds)
+
+
+            for interval in non_silent_intervals:
+                start, end = interval
                 segment_part = waveform[start:end]
-                segment_len = len(segment_part) / 16000
-                if segment_len < 1:
-                    continue
-                segments.append(segment_part)
-                start = end
 
-            for segment in segments:
+                # If adding the next part exceeds max duration, store the segment and start a new one
+                if current_length + len(segment_part) > max_duration:
+                    segments.append(np.concatenate(current_segment))
+                    current_segment = []
+                    current_length = 0
+
+                current_segment.append(segment_part)
+                current_length += len(segment_part)
+
+            # Append the last segment if it's not empty
+            if current_segment:
+                segments.append(np.concatenate(current_segment))
+
+            # Process each segment
+            outputs = []
+            for i, segment in enumerate(segments):
+                print(f"Processing segment {i + 1}/{len(segments)}, length: {len(segment) / sr:.2f} seconds")
+
+                # import soundfile as sf
+                # sf.write(f"outputs/segment_{i}.wav", segment, sr)
+
                 segment_tensor = torch.tensor(segment).to(device)
 
                 # Fake a batch for the segment
                 batch = segment_tensor.unsqueeze(0).to(device)
-                rel_length = torch.tensor([1.0]).to(device)
+                rel_length = torch.tensor([1.0]).to(device)  # Adjust if necessary
 
                 # Pass the segment through the ASR model
-                segment_output = self.encode_batch_whisper(device, batch, rel_length)
-                yield segment_output
+                outputs.append(self.encode_batch_w2v2(device, batch, rel_length))
+            return outputs
         else:
             waveform = torch.tensor(waveform).to(device)
             waveform = waveform.to(device)
+            # Fake a batch:
             batch = waveform.unsqueeze(0)
             rel_length = torch.tensor([1.0]).to(device)
-            outputs = self.encode_batch_whisper(device, batch, rel_length)
-            yield outputs
-
+            outputs = self.encode_batch_w2v2(device, batch, rel_length)
+            return [outputs]
 
     def classify_file_whisper(self, waveform, pipe, device):
         # waveform, sr = librosa.load(path, sr=16000)