141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
"""
|
|
/omr/pdf - PDF sheet music to MusicXML via Audiveris
|
|
/omr/image - Image (PNG/JPG) sheet music to MusicXML via Audiveris
|
|
"""
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from fastapi import APIRouter, HTTPException, UploadFile
|
|
from fastapi.responses import JSONResponse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
|
|
# C2 - file size cap
|
|
MAX_UPLOAD_BYTES = 50 * 1024 * 1024 # 50 MB
|
|
|
|
# H4 - extension allowlists
|
|
_PDF_EXTS = {".pdf"}
|
|
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tiff", ".tif"}
|
|
|
|
# M5 - magic bytes
|
|
_PDF_MAGIC = b"%PDF"
|
|
_PNG_MAGIC = b"\x89PNG"
|
|
_JPEG_MAGIC = b"\xff\xd8\xff"
|
|
_TIFF_MAGIC = (b"II*\x00", b"MM\x00*")
|
|
|
|
|
|
async def _read_with_limit(file: UploadFile, max_bytes: int = MAX_UPLOAD_BYTES) -> bytes:
|
|
"""Stream-read upload in 64 KB chunks; reject as soon as total exceeds limit.
|
|
|
|
Unlike ``await file.read()`` followed by a size check, this never buffers
|
|
an arbitrarily large payload into memory (H2: memory-DoS prevention).
|
|
"""
|
|
chunks: list[bytes] = []
|
|
total = 0
|
|
while True:
|
|
chunk = await file.read(65_536)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > max_bytes:
|
|
raise HTTPException(status_code=413, detail="Upload too large (max 50 MB)")
|
|
chunks.append(chunk)
|
|
return b"".join(chunks)
|
|
|
|
|
|
def _check_ext(filename: str, allowed: set[str]) -> None:
|
|
ext = Path(filename).suffix.lower()
|
|
if ext not in allowed:
|
|
raise HTTPException(
|
|
status_code=422,
|
|
detail=f"Unsupported file type. Allowed: {', '.join(sorted(allowed))}",
|
|
)
|
|
|
|
|
|
def _check_pdf_magic(data: bytes) -> None:
|
|
if not data.startswith(_PDF_MAGIC):
|
|
raise HTTPException(status_code=422, detail="Not a valid PDF file")
|
|
|
|
|
|
def _check_image_magic(data: bytes, ext: str) -> None:
|
|
if ext in {".png"} and not data.startswith(_PNG_MAGIC):
|
|
raise HTTPException(status_code=422, detail="Not a valid PNG file")
|
|
if ext in {".jpg", ".jpeg"} and not data.startswith(_JPEG_MAGIC):
|
|
raise HTTPException(status_code=422, detail="Not a valid JPEG file")
|
|
if ext in {".tiff", ".tif"} and not any(data.startswith(m) for m in _TIFF_MAGIC):
|
|
raise HTTPException(status_code=422, detail="Not a valid TIFF file")
|
|
|
|
|
|
@router.post("/pdf")
|
|
async def omr_pdf(file: UploadFile):
|
|
"""Run OMR on a PDF sheet music file to MusicXML."""
|
|
try:
|
|
_check_ext(file.filename or "score.pdf", _PDF_EXTS)
|
|
data = await _read_with_limit(file)
|
|
_check_pdf_magic(data)
|
|
|
|
from services.audiveris import pdf_to_musicxml
|
|
musicxml = await pdf_to_musicxml(data, suffix=".pdf")
|
|
return JSONResponse({"musicxml": musicxml}, headers={"Cache-Control": "no-store"})
|
|
except HTTPException:
|
|
raise
|
|
except Exception:
|
|
logger.exception("omr_pdf failed")
|
|
raise HTTPException(status_code=500, detail="OMR processing failed. Please try again.")
|
|
|
|
|
|
@router.post("/image")
|
|
async def omr_image(file: UploadFile):
|
|
"""Run OMR on an image (PNG/JPG) of sheet music to MusicXML."""
|
|
try:
|
|
filename = file.filename or "image.png"
|
|
_check_ext(filename, _IMAGE_EXTS)
|
|
data = await _read_with_limit(file)
|
|
ext = Path(filename).suffix.lower()
|
|
_check_image_magic(data, ext)
|
|
|
|
from services.audiveris import pdf_to_musicxml
|
|
musicxml = await pdf_to_musicxml(data, suffix=ext)
|
|
return JSONResponse({"musicxml": musicxml}, headers={"Cache-Control": "no-store"})
|
|
except HTTPException:
|
|
raise
|
|
except Exception:
|
|
logger.exception("omr_image failed")
|
|
raise HTTPException(status_code=500, detail="OMR processing failed. Please try again.")
|
|
|
|
|
|
@router.post("/pdf2midi")
|
|
async def omr_pdf_to_midi(file: UploadFile):
|
|
"""PDF sheet music, OMR to MusicXML + MIDI in one round-trip.
|
|
|
|
Returns JSON: { musicxml: str, midi_b64: str, filename: str }
|
|
Both artifacts arrive together so the frontend can render sheet music
|
|
AND enable playback without a second network request.
|
|
"""
|
|
import base64
|
|
|
|
try:
|
|
_check_ext(file.filename or "score.pdf", _PDF_EXTS)
|
|
data = await _read_with_limit(file)
|
|
_check_pdf_magic(data)
|
|
|
|
from services.audiveris import pdf_to_musicxml_and_midi
|
|
musicxml, midi_bytes = await pdf_to_musicxml_and_midi(data, suffix=".pdf")
|
|
midi_b64 = base64.b64encode(midi_bytes).decode("ascii")
|
|
|
|
return JSONResponse(
|
|
{"musicxml": musicxml, "midi_b64": midi_b64, "filename": "score.mid"},
|
|
headers={"Cache-Control": "no-store"},
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except RuntimeError as exc:
|
|
# User-safe messages from audiveris.py (no JAR, no output, etc.)
|
|
logger.warning("omr_pdf_to_midi: %s", exc)
|
|
raise HTTPException(status_code=422, detail=str(exc))
|
|
except Exception:
|
|
logger.exception("omr_pdf_to_midi failed unexpectedly")
|
|
raise HTTPException(status_code=500, detail="PDF processing failed. Please try again.")
|