Update app.py

app.py

@@ -1,59 +1,73 @@
 import gradio as gr
-import tempfile
-import requests
 from moviepy.editor import VideoFileClip
+from speechbrain.pretrained import EncoderClassifier
 import torchaudio
-import torch
-from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2FeatureExtractor, pipeline
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from pytube import YouTube
+import os
 
-# Load Models
-accent_model = Wav2Vec2ForSequenceClassification.from_pretrained("sreyan88/wav2vec2-large-xlsr-53-english-accent")
-feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("sreyan88/wav2vec2-large-xlsr-53-english-accent")
+CLASSIFIER = "Jzuluaga/accent-id-commonaccent_xlsr-en-english"
 
-whisper = pipeline("automatic-speech-recognition", model="openai/whisper-small")
+def download_video(url):
+    """Handles YouTube and direct video links"""
+    if "youtube.com" in url:
+        yt = YouTube(url)
+        stream = yt.streams.filter(progressive=True, file_extension='mp4').first()
+        video_path = stream.download()
+        return video_path
+    else:  # Direct download (assumes URL is direct mp4 link)
+        # Download file locally
+        import requests
+        local_filename = "temp_video.mp4"
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            with open(local_filename, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        return local_filename
 
-# Helpers
-def download_and_extract_audio(url):
-    r = requests.get(url, stream=True)
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as video_file:
-        for chunk in r.iter_content(chunk_size=1024):
-            if chunk:
-                video_file.write(chunk)
-        video_path = video_file.name
+def extract_audio(video_path):
     clip = VideoFileClip(video_path)
-    audio_path = video_path.replace(".mp4", ".wav")
-    clip.audio.write_audiofile(audio_path)
+    audio_path = "temp_audio.wav"
+    clip.audio.write_audiofile(audio_path, logger=None)
+    clip.close()
     return audio_path
 
 def classify_accent(audio_path):
-    # Load and process audio
+    classifier = EncoderClassifier.from_hparams(
+        source=CLASSIFIER,
+        savedir="pretrained_models/accent_classifier",
+        run_opts={"device":"cpu"}  # or "cuda" if GPU available
+    )
     waveform, sample_rate = torchaudio.load(audio_path)
-    inputs = feature_extractor(waveform[0], sampling_rate=sample_rate, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        logits = accent_model(**inputs).logits
-    predicted_class = torch.argmax(logits, dim=1).item()
-    confidence = torch.softmax(logits, dim=1)[0, predicted_class].item()
-    labels = ["american", "australian", "british", "indian", "canadian"]  # match model classes
-    return labels[predicted_class], round(confidence * 100, 2)
+    prediction = classifier.classify_batch(waveform)
+    # prediction format: (scores, probabilities, embeddings, predicted_labels)
+    predicted_accent = prediction[3][0]
+    confidence = prediction[1].exp().max().item() * 100
+    return predicted_accent, f"{confidence:.2f}%"
 
-def full_pipeline(url):
-    audio_path = download_and_extract_audio(url)
-    accent, confidence = classify_accent(audio_path)
-    transcription = whisper(audio_path)["text"]
-    summary = f"The speaker’s accent is most likely **{accent.capitalize()}** with a confidence score of **{confidence}%**."
-    return transcription, accent.capitalize(), confidence, summary
+def process_video(url):
+    try:
+        video_path = download_video(url)
+        audio_path = extract_audio(video_path)
+        accent, confidence = classify_accent(audio_path)
+    finally:
+        # Cleanup temp files if they exist
+        for f in [video_path, audio_path]:
+            if os.path.exists(f):
+                os.remove(f)
+    return accent, confidence
 
-# Gradio Interface
-gr.Interface(
-    fn=full_pipeline,
-    inputs=gr.Textbox(label="Public Video URL (MP4 or Loom)"),
+# Gradio interface
+iface = gr.Interface(
+    fn=process_video,
+    inputs=gr.Textbox(label="Enter Public Video URL (YouTube, Loom, direct MP4)"),
     outputs=[
-        gr.Textbox(label="Transcription"),
-        gr.Textbox(label="Accent"),
-        gr.Number(label="Confidence Score (%)"),
-        gr.Textbox(label="Summary"),
+        gr.Textbox(label="Detected Accent"),
+        gr.Textbox(label="Confidence Score")
     ],
-    title="Accent Classifier for English Speakers",
-    description="Paste a video URL to detect English accent and transcription. Built with Wav2Vec2 and Whisper on Hugging Face 🤗."
-).launch()
+    title="English Accent Classifier",
+    description="Paste a public video URL to detect the English accent and confidence score."
+)
+
+if __name__ == "__main__":
+    iface.launch()