diff --git a/pyproject.toml b/pyproject.toml index 8fbf919..83436a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ requires-python = ">=3.12" dependencies = [ "fastapi>=0.115", "uvicorn>=0.34", + "python-multipart>=0.0.17", "pdf2image>=1.17", "pytesseract>=0.3.13", "opencv-python>=4.10", @@ -18,6 +19,8 @@ dependencies = [ dev = [ "pytest>=8.0", "httpx>=0.28", + "anyio>=4.0", + "pytest-anyio>=0.0.0", ] [build-system] diff --git a/tests/test_app.py b/tests/test_app.py new file mode 100644 index 0000000..31ed700 --- /dev/null +++ b/tests/test_app.py @@ -0,0 +1,79 @@ +import io +from pathlib import Path + +import pytest +from httpx import ASGITransport, AsyncClient + +from unredact.app import app + + +@pytest.fixture +def pdf_bytes(sample_pdf: Path) -> bytes: + return sample_pdf.read_bytes() + + +@pytest.mark.anyio +async def test_upload_pdf(pdf_bytes: bytes): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post( + "/api/upload", + files={"file": ("test.pdf", pdf_bytes, "application/pdf")}, + ) + assert resp.status_code == 200 + data = resp.json() + assert "doc_id" in data + assert data["page_count"] > 0 + + +@pytest.mark.anyio +async def test_get_page_overlay(pdf_bytes: bytes): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + # Upload first + resp = await client.post( + "/api/upload", + files={"file": ("test.pdf", pdf_bytes, "application/pdf")}, + ) + doc_id = resp.json()["doc_id"] + + # Get page 1 overlay + resp = await client.get(f"/api/doc/{doc_id}/page/1/overlay") + assert resp.status_code == 200 + assert resp.headers["content-type"] == "image/png" + assert len(resp.content) > 1000 # Should be a real image + + +@pytest.mark.anyio +async def test_get_page_original(pdf_bytes: bytes): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post( + "/api/upload", + files={"file": ("test.pdf", pdf_bytes, "application/pdf")}, + ) + doc_id = resp.json()["doc_id"] + + resp = await client.get(f"/api/doc/{doc_id}/page/1/original") + assert resp.status_code == 200 + assert resp.headers["content-type"] == "image/png" + + +@pytest.mark.anyio +async def test_get_page_data(pdf_bytes: bytes): + """Should return OCR data and font info as JSON.""" + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.post( + "/api/upload", + files={"file": ("test.pdf", pdf_bytes, "application/pdf")}, + ) + doc_id = resp.json()["doc_id"] + + resp = await client.get(f"/api/doc/{doc_id}/page/1/data") + assert resp.status_code == 200 + data = resp.json() + assert "font" in data + assert "lines" in data + assert data["font"]["name"] + assert data["font"]["size"] > 0 diff --git a/unredact/app.py b/unredact/app.py new file mode 100644 index 0000000..e09ed3e --- /dev/null +++ b/unredact/app.py @@ -0,0 +1,115 @@ +import io +import uuid +from pathlib import Path +from tempfile import TemporaryDirectory + +from fastapi import FastAPI, UploadFile +from fastapi.responses import Response, JSONResponse +from fastapi.staticfiles import StaticFiles +from PIL import Image + +from unredact.pipeline.rasterize import rasterize_pdf +from unredact.pipeline.ocr import ocr_page +from unredact.pipeline.font_detect import detect_font +from unredact.pipeline.overlay import render_overlay + +app = FastAPI(title="Unredact") + +# In-memory store for uploaded docs (local-only tool, no persistence needed) +_docs: dict[str, dict] = {} + +STATIC_DIR = Path(__file__).parent / "static" +if STATIC_DIR.exists(): + app.mount("/static", StaticFiles(directory=str(STATIC_DIR), html=True), name="static") + + +def _image_to_png_bytes(img: Image.Image) -> bytes: + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +@app.post("/api/upload") +async def upload_pdf(file: UploadFile): + content = await file.read() + doc_id = uuid.uuid4().hex[:12] + + # Write to temp file for pdf2image + tmp = TemporaryDirectory() + pdf_path = Path(tmp.name) / "doc.pdf" + pdf_path.write_bytes(content) + + pages = rasterize_pdf(pdf_path) + + # Process each page + page_data = {} + for i, page_img in enumerate(pages, start=1): + lines = ocr_page(page_img) + font_match = detect_font(lines, page_img) + overlay_img = render_overlay(page_img, lines, font_match) + page_data[i] = { + "original": page_img, + "overlay": overlay_img, + "lines": lines, + "font_match": font_match, + } + + _docs[doc_id] = { + "page_count": len(pages), + "pages": page_data, + "tmp": tmp, # prevent cleanup + } + + return {"doc_id": doc_id, "page_count": len(pages)} + + +@app.get("/api/doc/{doc_id}/page/{page}/original") +async def get_page_original(doc_id: str, page: int): + doc = _docs.get(doc_id) + if not doc or page not in doc["pages"]: + return JSONResponse({"error": "not found"}, status_code=404) + png = _image_to_png_bytes(doc["pages"][page]["original"]) + return Response(content=png, media_type="image/png") + + +@app.get("/api/doc/{doc_id}/page/{page}/overlay") +async def get_page_overlay(doc_id: str, page: int): + doc = _docs.get(doc_id) + if not doc or page not in doc["pages"]: + return JSONResponse({"error": "not found"}, status_code=404) + overlay = doc["pages"][page]["overlay"] + # Convert RGBA to RGB for PNG output + rgb = overlay.convert("RGB") + png = _image_to_png_bytes(rgb) + return Response(content=png, media_type="image/png") + + +@app.get("/api/doc/{doc_id}/page/{page}/data") +async def get_page_data(doc_id: str, page: int): + doc = _docs.get(doc_id) + if not doc or page not in doc["pages"]: + return JSONResponse({"error": "not found"}, status_code=404) + + pd = doc["pages"][page] + fm = pd["font_match"] + lines_json = [] + for line in pd["lines"]: + chars_json = [ + {"text": c.text, "x": c.x, "y": c.y, "w": c.w, "h": c.h, "conf": c.conf} + for c in line.chars + ] + lines_json.append({ + "text": line.text, + "x": line.x, "y": line.y, "w": line.w, "h": line.h, + "chars": chars_json, + }) + + return { + "font": { + "name": fm.font_name, + "size": fm.font_size, + "score": fm.score, + "path": str(fm.font_path), + }, + "lines": lines_json, + }