Spaces:
No application file
No application file
| import io | |
| import os | |
| import json | |
| from pathlib import Path | |
| from fastapi import FastAPI, UploadFile, File, Form, HTTPException | |
| from fastapi.responses import StreamingResponse, JSONResponse | |
| from batch_processor import batch_embed, batch_extract | |
| from watermark import embed_watermark, extract_watermark | |
| app = FastAPI(title="AudioWatermarker API") | |
| async def api_embed( | |
| file: UploadFile = File(...), | |
| message: str = Form(...), | |
| password: str = Form(...) | |
| ): | |
| """Inserta watermark en un único archivo.""" | |
| if file.filename is None: | |
| raise HTTPException(status_code=400, detail="No filename provided") | |
| suffix = Path(file.filename).suffix | |
| data = await file.read() | |
| # Guardar temp | |
| tmp_in = f"/tmp/in{suffix}" | |
| tmp_out = f"/tmp/out{suffix}" | |
| with open(tmp_in, "wb") as f: f.write(data) | |
| bits = embed_watermark(tmp_in, tmp_out, message.encode(), password) | |
| def iterfile(): | |
| with open(tmp_out, "rb") as f: | |
| yield from f | |
| headers = {"X-Bits-Inserted": str(bits)} | |
| return StreamingResponse(iterfile(), media_type="audio/wav", headers=headers) | |
| async def api_extract( | |
| file: UploadFile = File(...), | |
| password: str = Form(...) | |
| ): | |
| """Extrae watermark de un archivo.""" | |
| if file.filename is None: | |
| raise HTTPException(status_code=400, detail="No filename provided") | |
| suffix = Path(file.filename).suffix | |
| tmp_in = f"/tmp/in{suffix}" | |
| with open(tmp_in, "wb") as f: f.write(await file.read()) | |
| try: | |
| payload = extract_watermark(tmp_in, password) | |
| return JSONResponse({"payload": payload.decode(errors="ignore")}) | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| async def api_batch_embed( | |
| message: str = Form(...), | |
| password: str = Form(...), | |
| files: list[UploadFile] = File(...) | |
| ): | |
| """Batch embed: recibe múltiples archivos, devuelve ZIP con watermarked.""" | |
| import zipfile, tempfile | |
| tmpdir = tempfile.mkdtemp() | |
| outdir = tempfile.mkdtemp() | |
| # guardar entrantes | |
| for f in files: | |
| if f.filename is None: | |
| continue | |
| path = os.path.join(tmpdir, f.filename) | |
| with open(path, "wb") as fh: fh.write(await f.read()) | |
| results = batch_embed(tmpdir, outdir, message.encode(), password) | |
| # crear zip | |
| zip_path = os.path.join(tempfile.gettempdir(), "batch_embed.zip") | |
| with zipfile.ZipFile(zip_path, "w") as z: | |
| for fn in os.listdir(outdir): | |
| z.write(os.path.join(outdir, fn), arcname=fn) | |
| # retornar | |
| return StreamingResponse( | |
| open(zip_path, "rb"), media_type="application/zip", | |
| headers={"X-Results": json.dumps(results)} | |
| ) | |
| async def api_batch_extract( | |
| password: str = Form(...), | |
| files: list[UploadFile] = File(...) | |
| ): | |
| """Batch extract: retorna JSON con payloads o errores.""" | |
| import tempfile, os | |
| tmpdir = tempfile.mkdtemp() | |
| for f in files: | |
| if f.filename is None: | |
| continue | |
| path = os.path.join(tmpdir, f.filename) | |
| with open(path, "wb") as fh: fh.write(await f.read()) | |
| results = batch_extract(tmpdir, password) | |
| return JSONResponse(results) |