import atexit import os import tempfile import uuid from pathlib import Path import torch from flask import Flask, jsonify, render_template_string, request, send_from_directory from classes import ESC50_CLASSES from preprocessing import preprocess_audio from utils import load_compressed_model, load_model, predict BASE_DIR = Path(__file__).resolve().parent ORIGINAL_MODEL_PATH = BASE_DIR / "weights" / "esc50_model.pth" COMPRESSED_MODEL_PATH = BASE_DIR / "weights" / "esc50_model_compressed.pth" STATS_PATH = BASE_DIR / "stats" / "esc50_mel_stats.json" SAMPLES_DIR = BASE_DIR / "samples" UPLOAD_DIR = Path(tempfile.gettempdir()) / "soundedge_uploads" UPLOAD_DIR.mkdir(parents=True, exist_ok=True) ALLOWED_EXTENSIONS = {".wav"} LOW_CONFIDENCE_THRESHOLD = 0.6 MAX_UPLOAD_BYTES = 5 * 1024 * 1024 app = Flask(__name__) app.config["MAX_CONTENT_LENGTH"] = MAX_UPLOAD_BYTES _gpu_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") _model_cache = {"Original": None, "Compressed": None} def _safe_label(raw_name: str) -> str: return raw_name.replace("_", " ").title() def _sample_files() -> list[str]: if not SAMPLES_DIR.is_dir(): return [] return sorted([f.name for f in SAMPLES_DIR.iterdir() if f.suffix.lower() == ".wav"]) def _validate_wav_filename(filename: str) -> bool: ext = Path(filename).suffix.lower() return ext in ALLOWED_EXTENSIONS def _get_active_model(model_choice: str): if model_choice == "Compressed": if _model_cache["Compressed"] is None: _model_cache["Compressed"] = load_compressed_model( str(ORIGINAL_MODEL_PATH), str(COMPRESSED_MODEL_PATH), num_classes=len(ESC50_CLASSES), ) return _model_cache["Compressed"], torch.device("cpu") if _model_cache["Original"] is None: _model_cache["Original"] = load_model( str(ORIGINAL_MODEL_PATH), _gpu_device, num_classes=len(ESC50_CLASSES), ) return _model_cache["Original"], _gpu_device def _prediction_payload(model_choice: str, source_path: Path): model, device = _get_active_model(model_choice) input_tensor = preprocess_audio(str(source_path), str(STATS_PATH)) top_class, top_prob, all_probs = predict(model, input_tensor, device) return { "topClass": top_class, "topClassLabel": _safe_label(top_class), "topProbability": top_prob, "topProbabilityPct": round(top_prob * 100, 2), "lowConfidence": top_prob < LOW_CONFIDENCE_THRESHOLD, "top3": [ { "className": p["class_name"], "classLabel": _safe_label(p["class_name"]), "probability": p["probability"], "probabilityPct": round(p["probability"] * 100, 2), } for p in all_probs[:3] ], "allProbs": [ { "className": p["class_name"], "classLabel": _safe_label(p["class_name"]), "probability": p["probability"], "probabilityPct": round(p["probability"] * 100, 2), } for p in all_probs ], } @app.get("/") def index(): classes = [_safe_label(c) for c in ESC50_CLASSES] samples = [{"name": s, "label": _safe_label(Path(s).stem)} for s in _sample_files()] return render_template_string( PAGE_TEMPLATE, class_badges=classes, samples=samples, ) @app.get("/samples/") def serve_sample(filename: str): return send_from_directory(SAMPLES_DIR, filename) @app.get("/uploads/") def serve_uploaded_file(filename: str): return send_from_directory(UPLOAD_DIR, filename) @app.post("/upload") def upload_audio(): uploaded = request.files.get("file") if uploaded is None: return jsonify({"error": "Missing file field. Use multipart form key 'file'."}), 400 filename = uploaded.filename or "" if not filename: return jsonify({"error": "No file selected."}), 400 if not _validate_wav_filename(filename): return jsonify({"error": "Only .wav files are supported."}), 400 file_id = uuid.uuid4().hex stored_name = f"{file_id}.wav" destination = UPLOAD_DIR / stored_name uploaded.save(destination) return jsonify( { "fileId": file_id, "filename": filename, "audioUrl": f"/uploads/{stored_name}", } ) @app.post("/predict") def predict_audio(): payload = request.get_json(silent=True) or {} source_type = payload.get("source", "upload") model_choice = payload.get("model", "Original") if model_choice not in ("Original", "Compressed"): return jsonify({"error": "Invalid model. Use 'Original' or 'Compressed'."}), 400 try: if source_type == "upload": file_id = payload.get("fileId", "") if not file_id: return jsonify({"error": "Missing fileId for uploaded source."}), 400 source_path = UPLOAD_DIR / f"{file_id}.wav" if not source_path.exists(): return jsonify({"error": "Uploaded file not found. Upload again."}), 404 elif source_type == "sample": sample_name = payload.get("sampleName", "") if not sample_name: return jsonify({"error": "Missing sampleName for sample source."}), 400 source_path = SAMPLES_DIR / sample_name if not source_path.exists() or source_path.suffix.lower() != ".wav": return jsonify({"error": "Sample not found."}), 404 else: return jsonify({"error": "Invalid source. Use 'upload' or 'sample'."}), 400 result = _prediction_payload(model_choice=model_choice, source_path=source_path) return jsonify(result) except Exception as exc: return jsonify({"error": f"Error during inference: {exc}"}), 500 @app.get("/health") def health(): return jsonify({"status": "ok"}) def _cleanup_uploads(): if not UPLOAD_DIR.exists(): return for wav in UPLOAD_DIR.glob("*.wav"): try: wav.unlink() except OSError: pass atexit.register(_cleanup_uploads) PAGE_TEMPLATE = """ SoundEdge - Environmental Sound Classification

🔊SoundEdge

Environmental Sound Classification - upload a short audio clip and let the model identify the sound.

Supported Sound Classes
{% for cls in class_badges %} {{ cls }} {% endfor %}
Select Model
Choose Audio Input
Upload Guide
Format: WAV (.wav) only.
Duration: Around 5 seconds gives best results.
Size: Keep under 5 MB.
{% if samples %}
{% else %}
No sample files found in the samples folder.
{% endif %}
""" if __name__ == "__main__": port = int(os.environ.get("PORT", "8501")) host = os.environ.get("HOST", "127.0.0.1") app.run(host=host, port=port, debug=False)