Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py - Musical Instrument Detection API con modelo específico
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import pipeline
@@ -11,7 +10,6 @@ from datetime import datetime
 import torch
 from contextlib import asynccontextmanager
 import subprocess
-import wave
 
 # Configurar cache
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
@@ -30,7 +28,7 @@ logger = logging.getLogger(__name__)
 classifier = None
 
 async def load_model():
-    """Cargar modelo específico para instrumentos musicales"""
+    """Cargar modelo MIT/AST más balanceado"""
     global classifier
     try:
         logger.info("Iniciando carga del modelo...")
@@ -40,8 +38,8 @@ async def load_model():
         os.makedirs('/tmp/huggingface', exist_ok=True)
         os.makedirs('/tmp/numba_cache', exist_ok=True)
         
-        # MODELO ESPECÍFICO PARA INSTRUMENTOS MUSICALES
-        model_name = "dima806/musical_instrument_detection"
+        # MODELO MIT/AST - MÁS BALANCEADO (527 clases de AudioSet)
+        model_name = "MIT/ast-finetuned-audioset-10-10-0.4593"
         
         logger.info(f"Cargando modelo: {model_name}")
         
@@ -52,11 +50,28 @@ async def load_model():
             return_all_scores=True
         )
         
-        logger.info("Modelo de instrumentos musicales cargado exitosamente")
+        logger.info("Modelo MIT/AST cargado exitosamente (527 clases)")
         
     except Exception as e:
-        logger.error(f"Error cargando modelo: {e}")
-        classifier = None
+        logger.error(f"Error cargando modelo MIT/AST: {e}")
+        
+        # Fallback al modelo anterior si MIT/AST no funciona
+        try:
+            logger.info("Intentando modelo alternativo...")
+            model_name_fallback = "dima806/musical_instrument_detection"
+            
+            classifier = pipeline(
+                "audio-classification",
+                model=model_name_fallback,
+                device=-1,
+                return_all_scores=True
+            )
+            
+            logger.info("Modelo alternativo cargado")
+            
+        except Exception as e2:
+            logger.error(f"Error con modelo alternativo: {e2}")
+            classifier = None
 
 async def cleanup_model():
     """Limpiar recursos"""
@@ -72,8 +87,8 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="Musical Instrument Detection API",
-    description="API para detectar instrumentos musicales en pistas de audio",
-    version="3.0.0",
+    description="API para detectar instrumentos musicales con modelo balanceado",
+    version="4.0.0",
     lifespan=lifespan
 )
 
@@ -159,14 +174,63 @@ def load_audio_robust(file_path):
         logger.error(f"Error cargando audio: {e}")
         raise
 
+def filter_musical_instruments(results):
+    """Filtrar solo instrumentos musicales de los resultados"""
+    # Palabras clave de instrumentos musicales en AudioSet
+    instrument_keywords = [
+        'guitar', 'piano', 'drum', 'violin', 'flute', 'trumpet', 
+        'saxophone', 'bass', 'keyboard', 'harp', 'organ', 'cello',
+        'clarinet', 'trombone', 'harmonica', 'accordion', 'banjo',
+        'mandolin', 'ukulele', 'cymbal', 'tambourine', 'xylophone',
+        'marimba', 'vibraphone', 'synthesizer', 'electric', 'acoustic'
+    ]
+    
+    # Instrumentos específicos de AudioSet
+    musical_labels = [
+        'Guitar', 'Electric guitar', 'Acoustic guitar', 'Bass guitar',
+        'Piano', 'Electric piano', 'Keyboard (musical)',
+        'Drum', 'Drum kit', 'Snare drum', 'Bass drum', 'Cymbal',
+        'Violin, fiddle', 'Cello', 'Double bass',
+        'Flute', 'Clarinet', 'Saxophone', 'Trumpet', 'Trombone',
+        'Harmonica', 'Accordion', 'Organ', 'Harp',
+        'Synthesizer', 'Musical instrument', 'Plucked string instrument',
+        'Bowed string instrument', 'Wind instrument', 'Percussion'
+    ]
+    
+    filtered_results = []
+    
+    for result in results:
+        label = result['label']
+        confidence = float(result['score'])
+        
+        # Verificar si es un instrumento musical
+        is_instrument = (
+            any(keyword.lower() in label.lower() for keyword in instrument_keywords) or
+            label in musical_labels or
+            confidence > 0.3  # Incluir resultados con alta confianza
+        )
+        
+        if is_instrument and confidence > 0.05:  # Umbral mínimo 5%
+            # Limpiar nombre del instrumento
+            clean_name = label.replace('Sound_', '').replace('_', ' ').title()
+            
+            filtered_results.append({
+                "instrument": clean_name,
+                "confidence": round(confidence, 4),
+                "percentage": round(confidence * 100, 2),
+                "original_label": label
+            })
+    
+    return filtered_results
+
 @app.get("/")
 async def root():
     return {
         "message": "Musical Instrument Detection API",
         "status": "online",
-        "version": "3.0.0",
-        "model": "dima806/musical_instrument_detection",
-        "capabilities": ["guitar", "piano", "drums", "violin", "bass", "flute", "trumpet", "saxophone"],
+        "version": "4.0.0",
+        "model": "MIT/ast-finetuned-audioset-10-10-0.4593",
+        "features": ["Modelo balanceado", "527 clases de AudioSet", "Menos sesgo hacia guitarra"],
         "endpoints": {
             "health": "/health",
             "detect": "/detect",
@@ -179,15 +243,14 @@ async def health_check():
     return {
         "status": "online" if classifier is not None else "error",
         "model_loaded": classifier is not None,
-        "model_name": "dima806/musical_instrument_detection",
-        "timestamp": datetime.now().isoformat(),
-        "ffmpeg_available": subprocess.run(['ffmpeg', '-version'], 
-                                         capture_output=True).returncode == 0
+        "model_name": "MIT/ast-finetuned-audioset-10-10-0.4593",
+        "total_classes": 527,
+        "timestamp": datetime.now().isoformat()
     }
 
 @app.post("/detect")
 async def detect_instrument(audio: UploadFile = File(...)):
-    """Detectar instrumentos musicales en pista de audio"""
+    """Detectar instrumentos musicales con modelo balanceado"""
     start_time = datetime.now()
     
     try:
@@ -196,17 +259,16 @@ async def detect_instrument(audio: UploadFile = File(...)):
             raise HTTPException(status_code=503, detail="Modelo no disponible")
         
         logger.info(f"Procesando pista: {audio.filename}")
-        logger.info(f"Content-Type: {audio.content_type}")
         
         # Leer contenido
         content = await audio.read()
         file_size_mb = len(content) / (1024 * 1024)
         logger.info(f"Tamaño: {file_size_mb:.2f}MB")
         
-        if file_size_mb > 15:  # Aumentamos límite para pistas completas
-            raise HTTPException(status_code=413, detail="Archivo demasiado grande. Máximo 15MB")
+        if file_size_mb > 15:
+            raise HTTPException(status_code=413, detail="Archivo demasiado grande")
         
-        # Crear archivo temporal
+        # Procesar audio
         with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as temp_file:
             temp_file.write(content)
             temp_path = temp_file.name
@@ -221,34 +283,25 @@ async def detect_instrument(audio: UploadFile = File(...)):
             logger.info(f"Audio cargado: {duration:.2f}s, {sample_rate}Hz")
             
             if duration < 0.5:
-                raise HTTPException(status_code=400, detail="Audio demasiado corto para detectar instrumentos")
+                raise HTTPException(status_code=400, detail="Audio demasiado corto")
             
-            # Para pistas largas, analizar en segmentos
+            # Para pistas largas, analizar segmento representativo
             if duration > 60:
-                # Analizar primeros 60 segundos
                 max_samples = 60 * sample_rate
                 audio_data = audio_data[:max_samples]
-                logger.info("Analizando primeros 60 segundos de la pista")
+                logger.info("Analizando primeros 60 segundos")
             
             # Normalizar
             if np.max(np.abs(audio_data)) > 0:
                 audio_data = audio_data / np.max(np.abs(audio_data))
             
-            logger.info("Detectando instrumentos musicales...")
+            logger.info("Clasificando con modelo MIT/AST...")
             
-            # Clasificar instrumentos
+            # Clasificar con modelo balanceado
             results = classifier(audio_data)
             
-            # Procesar y filtrar resultados
-            instruments_detected = []
-            for result in results:
-                confidence = float(result['score'])
-                if confidence > 0.1:  # Solo instrumentos con confianza > 10%
-                    instruments_detected.append({
-                        "instrument": result['label'],
-                        "confidence": round(confidence, 4),
-                        "percentage": round(confidence * 100, 2)
-                    })
+            # Filtrar solo instrumentos musicales
+            instruments_detected = filter_musical_instruments(results)
             
             # Ordenar por confianza
             instruments_detected.sort(key=lambda x: x['confidence'], reverse=True)
@@ -257,14 +310,18 @@ async def detect_instrument(audio: UploadFile = File(...)):
             
             response = {
                 "success": True,
-                "instruments_detected": instruments_detected[:10],  # Top 10 instrumentos
+                "instruments_detected": instruments_detected[:8],  # Top 8 instrumentos
                 "total_instruments_found": len(instruments_detected),
+                "model_info": {
+                    "name": "MIT/ast-finetuned-audioset",
+                    "total_classes": 527,
+                    "description": "Modelo balanceado entrenado en AudioSet"
+                },
                 "audio_info": {
                     "filename": audio.filename,
                     "duration": round(duration, 2),
                     "size_mb": round(file_size_mb, 2),
-                    "sample_rate": sample_rate,
-                    "analyzed_duration": min(duration, 60)
+                    "sample_rate": sample_rate
                 },
                 "processing_time": round(processing_time, 3),
                 "timestamp": datetime.now().isoformat()