Reimplement

hf_app.py

@@ -1,7 +1,359 @@
-from fastapi import FastAPI
-from fastapi.middleware.wsgi import WSGIMiddleware
+from pathlib import Path
+from typing import Any, Dict, List
+
+import json
+import os
+import shutil
+
+import torch
+from fastapi import FastAPI, File, Form, Request, UploadFile
+from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from loguru import logger
+from werkzeug.utils import secure_filename
+
+import main as extractor
 
-from app import app as flask_app
 
 app = FastAPI()
-app.mount("/", WSGIMiddleware(flask_app))
+
+
+# Static files and templates -------------------------------------------------
+app.mount("/static", StaticFiles(directory="static"), name="static")
+templates = Jinja2Templates(directory="templates")
+
+
+def flask_like_url_for(endpoint: str, **kwargs: Any) -> str:
+    """Minimal Flask-like url_for for templates using filename= for static.
+
+    The Jinja template calls url_for('static', filename='css/styles.css'),
+    which is Flask style. We emulate that here so templates work unchanged.
+    """
+
+    if endpoint == "static":
+        filename = str(kwargs.get("filename", ""))
+        return "/static/" + filename.lstrip("/")
+
+    # Fallback: just return "/<endpoint>"; templates only use static.
+    return "/" + endpoint.lstrip("/")
+
+
+templates.env.globals["url_for"] = flask_like_url_for
+
+
+# Configuration -------------------------------------------------------------
+UPLOAD_FOLDER = Path("./uploads")
+OUTPUT_FOLDER = Path("./output")
+MAX_CONTENT_LENGTH = 500 * 1024 * 1024  # 500MB
+
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(OUTPUT_FOLDER, exist_ok=True)
+
+
+# Global model cache --------------------------------------------------------
+_model: Any = None
+
+
+def get_device_info() -> Dict[str, Any]:
+    """Get information about GPU/CPU availability."""
+
+    cuda_available = torch.cuda.is_available()
+    device = "cuda" if cuda_available else "cpu"
+
+    info: Dict[str, Any] = {
+        "device": device,
+        "cuda_available": cuda_available,
+        "device_name": None,
+        "device_count": 0,
+    }
+
+    if cuda_available:
+        info["device_name"] = torch.cuda.get_device_name(0)
+        info["device_count"] = torch.cuda.device_count()
+
+    return info
+
+
+def load_model_once() -> Any:
+    """Load the DocLayout-YOLO model once and cache it in this process."""
+
+    global _model
+    if _model is None:
+        logger.info("Loading DocLayout-YOLO model...")
+        _model = extractor.get_model()
+        logger.info("Model loaded successfully")
+    return _model
+
+
+# Routes --------------------------------------------------------------------
+@app.get("/", response_class=HTMLResponse)
+async def index(request: Request) -> HTMLResponse:
+    """Main page, equivalent to the Flask index route."""
+
+    device_info = get_device_info()
+    return templates.TemplateResponse(
+        "index.html", {"request": request, "device_info": device_info}
+    )
+
+
+@app.get("/api/device-info")
+async def device_info() -> Dict[str, Any]:
+    """API endpoint to get device information."""
+
+    return get_device_info()
+
+
+@app.post("/api/upload")
+async def upload_files(
+    request: Request,
+    files: List[UploadFile] = File(..., alias="files[]"),
+    extraction_mode: str = Form("images"),
+) -> JSONResponse:
+    """Handle multiple PDF file uploads (FastAPI version of Flask route)."""
+
+    if not files or all((f.filename or "") == "" for f in files):
+        return JSONResponse({"error": "No files selected"}, status_code=400)
+
+    include_images = extraction_mode != "markdown"
+    include_markdown = extraction_mode != "images"
+
+    results: List[Dict[str, Any]] = []
+
+    for upload in files:
+        filename = upload.filename or ""
+        if not filename.endswith(".pdf"):
+            continue
+
+        try:
+            safe_name = secure_filename(filename)
+            stem = Path(safe_name).stem
+
+            upload_path = UPLOAD_FOLDER / safe_name
+            # Save uploaded file to disk
+            with upload_path.open("wb") as out_f:
+                while True:
+                    chunk = await upload.read(1024 * 1024)
+                    if not chunk:
+                        break
+                    out_f.write(chunk)
+
+            # Prepare output directory
+            output_dir = OUTPUT_FOLDER / stem
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            # Move PDF into output directory
+            pdf_path = output_dir / safe_name
+            upload_path.replace(pdf_path)
+
+            # Process PDF
+            extractor.USE_MULTIPROCESSING = False
+            logger.info(
+                f"Processing {safe_name} (images={include_images}, markdown={include_markdown})"
+            )
+
+            if include_images:
+                load_model_once()
+
+            extractor.process_pdf_with_pool(
+                pdf_path,
+                output_dir,
+                pool=None,
+                extract_images=include_images,
+                extract_markdown=include_markdown,
+            )
+
+            # Collect results
+            json_path = output_dir / f"{stem}_content_list.json"
+            elements: List[Dict[str, Any]] = []
+            if include_images and json_path.exists():
+                elements = json.loads(json_path.read_text(encoding="utf-8"))
+
+            annotated_pdf: str | None = None
+            if include_images:
+                candidate_pdf = output_dir / f"{stem}_layout.pdf"
+                if candidate_pdf.exists():
+                    annotated_pdf = str(candidate_pdf.relative_to(OUTPUT_FOLDER))
+
+            markdown_path: str | None = None
+            if include_markdown:
+                candidate_md = output_dir / f"{stem}.md"
+                if candidate_md.exists():
+                    markdown_path = str(candidate_md.relative_to(OUTPUT_FOLDER))
+
+            figures = [e for e in elements if e.get("type") == "figure"]
+            tables = [e for e in elements if e.get("type") == "table"]
+
+            results.append(
+                {
+                    "filename": safe_name,
+                    "stem": stem,
+                    "output_dir": str(output_dir.relative_to(OUTPUT_FOLDER)),
+                    "figures_count": len(figures),
+                    "tables_count": len(tables),
+                    "elements_count": len(elements),
+                    "annotated_pdf": annotated_pdf,
+                    "markdown_path": markdown_path,
+                    "include_images": include_images,
+                    "include_markdown": include_markdown,
+                }
+            )
+
+        except Exception as e:  # pragma: no cover - runtime error path
+            logger.error(f"Error processing {filename}: {e}")
+            results.append({"filename": filename, "error": str(e)})
+
+    return JSONResponse({"results": results})
+
+
+@app.get("/api/pdf-list")
+async def pdf_list() -> Dict[str, Any]:
+    """Get list of processed PDFs."""
+
+    pdfs: List[Dict[str, Any]] = []
+    output_dir = OUTPUT_FOLDER
+
+    if not output_dir.exists():
+        return {"pdfs": pdfs}
+
+    for item in output_dir.iterdir():
+        if item.is_dir():
+            json_files = list(item.glob("*_content_list.json"))
+            md_files = list(item.glob("*.md"))
+            pdf_files = list(item.glob("*.pdf"))
+
+            if json_files or md_files or pdf_files:
+                stem = item.name
+                pdfs.append(
+                    {
+                        "stem": stem,
+                        "output_dir": str(item.relative_to(output_dir)),
+                    }
+                )
+
+    return {"pdfs": pdfs}
+
+
+@app.get("/api/pdf-details/{pdf_stem:path}")
+async def pdf_details(pdf_stem: str) -> JSONResponse:
+    """Get detailed information about a processed PDF."""
+
+    output_dir = OUTPUT_FOLDER / pdf_stem
+
+    if not output_dir.exists():
+        return JSONResponse({"error": "PDF not found"}, status_code=404)
+
+    json_files = list(output_dir.glob("*_content_list.json"))
+    elements: List[Dict[str, Any]] = []
+    if json_files:
+        elements = json.loads(json_files[0].read_text(encoding="utf-8"))
+
+    figures = [e for e in elements if e.get("type") == "figure"]
+    tables = [e for e in elements if e.get("type") == "table"]
+
+    annotated_pdf: str | None = None
+    pdf_files = list(output_dir.glob("*_layout.pdf"))
+    if pdf_files:
+        annotated_pdf = str(pdf_files[0].relative_to(OUTPUT_FOLDER))
+
+    markdown_path: str | None = None
+    md_files = list(output_dir.glob("*.md"))
+    if md_files:
+        markdown_path = str(md_files[0].relative_to(OUTPUT_FOLDER))
+
+    figure_dir = output_dir / "figures"
+    table_dir = output_dir / "tables"
+
+    figure_images: List[str] = []
+    if figure_dir.exists():
+        figure_images = [
+            str(f.relative_to(OUTPUT_FOLDER)) for f in sorted(figure_dir.glob("*.png"))
+        ]
+
+    table_images: List[str] = []
+    if table_dir.exists():
+        table_images = [
+            str(t.relative_to(OUTPUT_FOLDER)) for t in sorted(table_dir.glob("*.png"))
+        ]
+
+    return JSONResponse(
+        {
+            "stem": pdf_stem,
+            "figures": figures,
+            "tables": tables,
+            "figures_count": len(figures),
+            "tables_count": len(tables),
+            "elements_count": len(elements),
+            "annotated_pdf": annotated_pdf,
+            "markdown_path": markdown_path,
+            "figure_images": figure_images,
+            "table_images": table_images,
+        }
+    )
+
+
+@app.get("/output/{filename:path}")
+async def output_file(filename: str) -> FileResponse | JSONResponse:
+    """Serve output files (PDFs, images, markdown)."""
+
+    output_root = OUTPUT_FOLDER.resolve()
+    file_path = (output_root / filename).resolve()
+
+    if output_root not in file_path.parents and file_path != output_root:
+        return JSONResponse({"error": "Invalid path"}, status_code=400)
+
+    if not file_path.exists() or not file_path.is_file():
+        return JSONResponse({"error": "Not found"}, status_code=404)
+
+    return FileResponse(file_path)
+
+
+def _delete_by_stem(stem_raw: str) -> JSONResponse:
+    stem = (stem_raw or "").strip()
+    if not stem:
+        return JSONResponse({"error": "Missing stem"}, status_code=400)
+
+    output_root = OUTPUT_FOLDER.resolve()
+    target_dir = (output_root / stem).resolve()
+
+    if output_root not in target_dir.parents and target_dir != output_root:
+        return JSONResponse({"error": "Invalid stem path"}, status_code=400)
+
+    if not target_dir.exists() or not target_dir.is_dir():
+        return JSONResponse({"error": "Not found"}, status_code=404)
+
+    shutil.rmtree(target_dir, ignore_errors=False)
+    logger.info(f"Deleted processed output: {target_dir}")
+
+    return JSONResponse({"ok": True, "deleted": stem})
+
+
+@app.post("/api/delete")
+async def delete_pdf(request: Request, stem_form: str | None = Form(default=None)) -> JSONResponse:
+    """Delete a processed PDF directory by stem (JSON or form body)."""
+
+    try:
+        stem = (stem_form or "").strip()
+        if not stem:
+            data: Dict[str, Any] = {}
+            try:
+                data = await request.json()
+            except Exception:
+                data = {}
+            stem = (str(data.get("stem") or "")).strip()
+        return _delete_by_stem(stem)
+    except Exception as e:  # pragma: no cover - runtime error path
+        logger.error(f"Delete failed: {e}")
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+
+@app.api_route("/api/delete/{stem:path}", methods=["POST", "GET"])
+async def delete_pdf_by_path(stem: str) -> JSONResponse:
+    """Alternate endpoint to delete using URL path, for clients avoiding bodies."""
+
+    try:
+        return _delete_by_stem(stem)
+    except Exception as e:  # pragma: no cover - runtime error path
+        logger.error(f"Delete failed: {e}")
+        return JSONResponse({"error": str(e)}, status_code=500)
+