app.py

"""
Baby Cry Analysis - Flask API for Hugging Face Spaces
Exposes REST endpoint for audio-based baby cry classification.
Wrapped with Gradio to keep the Space alive.
"""

from flask import Flask, request, jsonify
import tempfile
import os
import logging
import threading

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import inference module (model loads at import time)
from inference import analyze_cry, preload_model

# Preload model at startup
preload_model()

app = Flask(__name__)

# Maximum file size: 16MB
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024

# Allowed audio extensions
ALLOWED_EXTENSIONS = {'wav', 'mp3', 'ogg', 'flac', 'm4a', 'webm'}


def allowed_file(filename: str) -> bool:
    """Check if file extension is allowed."""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route("/health", methods=["GET"])
def health():
    """Health check endpoint."""
    return jsonify({"status": "healthy", "service": "baby-cry-ai"})


@app.route("/analyze-cry", methods=["POST"])
def analyze():
    """
    Analyze baby cry audio file.
    
    Expects multipart/form-data with 'audio' file field.
    
    Returns:
        JSON with cry analysis results:
        - cry_detected: boolean
        - top_reason: string (hunger, belly_pain, tired, discomfort, burping)
        - scores: object with confidence scores per label
        - disclaimer: legal disclaimer string
    """
    # Validate request has audio file
    if "audio" not in request.files:
        return jsonify({
            "error": "No audio file provided",
            "message": "Please upload an audio file with key 'audio'"
        }), 400
    
    audio = request.files["audio"]
    
    # Validate filename exists
    if audio.filename == '':
        return jsonify({
            "error": "Empty filename",
            "message": "No file selected"
        }), 400
    
    # Validate file extension
    if not allowed_file(audio.filename):
        return jsonify({
            "error": "Invalid file type",
            "message": f"Allowed types: {', '.join(ALLOWED_EXTENSIONS)}"
        }), 400
    
    # Get original extension for temp file
    ext = audio.filename.rsplit('.', 1)[1].lower()
    
    try:
        # Save to temporary file for processing
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") as tmp:
            audio.save(tmp.name)
            temp_path = tmp.name
        
        logger.info(f"Processing audio file: {temp_path}")
        
        # Run inference
        result = analyze_cry(temp_path)
        
        # Add disclaimer
        result["disclaimer"] = "AI-generated suggestion, not a medical diagnosis. Please consult a healthcare professional for medical advice."
        
        logger.info(f"Analysis successful: {result['top_reason']}")
        return jsonify(result)
    
    except Exception as e:
        logger.error(f"Analysis failed: {str(e)}", exc_info=True)
        return jsonify({
            "error": "Analysis failed",
            "message": str(e)
        }), 500
    
    finally:
        # Clean up temp file
        if 'temp_path' in locals() and os.path.exists(temp_path):
            os.unlink(temp_path)


@app.errorhandler(413)
def too_large(e):
    """Handle file too large error."""
    return jsonify({
        "error": "File too large",
        "message": "Maximum file size is 16MB"
    }), 413


def run_flask():
    """Run Flask app in background thread."""
    app.run(host="0.0.0.0", port=7860, debug=False, use_reloader=False)


# Start Flask in background thread
flask_thread = threading.Thread(target=run_flask, daemon=True)
flask_thread.start()

# Gradio UI to keep the Space alive
import gradio as gr

with gr.Blocks() as demo:
    gr.Markdown("""
    # 👶 Baby Cry Analysis API
    
    This Space hosts a REST API for analyzing baby cries using machine learning.
    
    ## 🔗 API Endpoints
    
    ### Health Check
    ```
    GET /health
    ```
    
    ### Analyze Cry
    ```
    POST /analyze-cry
    Content-Type: multipart/form-data
    Body: audio=<file>
    ```
    
    **Supported formats:** WAV, MP3, OGG, FLAC, M4A, WebM
    
    ## 📡 Example Usage
    
    ```python
    import requests
    
    with open("baby_cry.wav", "rb") as f:
        response = requests.post(
            "https://YOUR-SPACE.hf.space/analyze-cry",
            files={"audio": f}
        )
    print(response.json())
    ```
    
    ## 📊 Response Format
    
    ```json
    {
      "cry_detected": true,
      "top_reason": "hunger",
      "scores": {
        "hunger": 0.45,
        "belly_pain": 0.20,
        "tired": 0.15,
        "discomfort": 0.12,
        "burping": 0.08
      },
      "disclaimer": "AI-generated suggestion..."
    }
    ```
    
    ---
    ⚠️ *This is an AI-generated suggestion, not a medical diagnosis.*
    """)

demo.launch(server_name="0.0.0.0", server_port=7861)