5ea18d76d6
Adds a GET /api/jobs/suggestions endpoint that returns distinct values for author, chapter, and reviewer_name from the database, and wires them into HTML datalist elements on the New Job, result view, and Browse Jobs pages. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
882 lines
31 KiB
Python
882 lines
31 KiB
Python
import os
|
|
import re
|
|
import uuid
|
|
import tempfile
|
|
import shutil
|
|
import base64
|
|
from typing import List, Dict, Any, Optional
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Query
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import JSONResponse, StreamingResponse, FileResponse
|
|
from pydantic import BaseModel
|
|
import torch
|
|
from transformers import AutoModel, AutoTokenizer
|
|
from PIL import Image
|
|
import uvicorn
|
|
from decouple import config as env_config
|
|
|
|
# Import PDF and document conversion utilities
|
|
from pdf_utils import (
|
|
pdf_to_images_high_quality,
|
|
images_to_pdf,
|
|
extract_ref_patterns,
|
|
crop_images_from_refs,
|
|
clean_markdown_content
|
|
)
|
|
from format_converter import DocumentConverter
|
|
from database import init_db, get_db
|
|
|
|
OCR_IMAGES_DIR = env_config("OCR_IMAGES_DIR", default="/data/ocr_images")
|
|
|
|
# -----------------------------
|
|
# Lifespan context for model loading
|
|
# -----------------------------
|
|
model = None
|
|
tokenizer = None
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Load model on startup, cleanup on shutdown"""
|
|
global model, tokenizer
|
|
|
|
# Image storage directory
|
|
os.makedirs(OCR_IMAGES_DIR, exist_ok=True)
|
|
|
|
# Database
|
|
try:
|
|
init_db()
|
|
except Exception as exc:
|
|
print(f"Warning: database initialization failed: {exc}")
|
|
|
|
# Environment setup
|
|
os.environ.pop("TRANSFORMERS_CACHE", None)
|
|
MODEL_NAME = env_config("MODEL_NAME", default="deepseek-ai/DeepSeek-OCR")
|
|
HF_HOME = env_config("HF_HOME", default="/models")
|
|
os.makedirs(HF_HOME, exist_ok=True)
|
|
|
|
# Load model
|
|
print(f"🚀 Loading {MODEL_NAME}...")
|
|
torch_dtype = torch.bfloat16
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
MODEL_NAME,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
model = AutoModel.from_pretrained(
|
|
MODEL_NAME,
|
|
trust_remote_code=True,
|
|
use_safetensors=True,
|
|
attn_implementation="eager",
|
|
torch_dtype=torch_dtype,
|
|
).eval().to("cuda")
|
|
|
|
# Pad token setup
|
|
try:
|
|
if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token_id", None) is not None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
if getattr(model.config, "pad_token_id", None) is None and getattr(tokenizer, "pad_token_id", None) is not None:
|
|
model.config.pad_token_id = tokenizer.pad_token_id
|
|
except Exception:
|
|
pass
|
|
|
|
print("✅ Model loaded and ready!")
|
|
|
|
yield
|
|
|
|
# Cleanup
|
|
print("🛑 Shutting down...")
|
|
|
|
# -----------------------------
|
|
# FastAPI app
|
|
# -----------------------------
|
|
app = FastAPI(
|
|
title="DeepSeek-OCR API",
|
|
description="Blazing fast OCR with DeepSeek-OCR model 🔥",
|
|
version="2.0.0",
|
|
lifespan=lifespan
|
|
)
|
|
|
|
# CORS middleware for React frontend
|
|
CORS_ORIGINS = env_config("CORS_ORIGINS", default="").split(",")
|
|
CORS_ORIGINS = [o.strip() for o in CORS_ORIGINS if o.strip()]
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=CORS_ORIGINS if CORS_ORIGINS else ["http://localhost:3000"],
|
|
allow_credentials=True,
|
|
allow_methods=["GET", "POST"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# -----------------------------
|
|
# Prompt builder
|
|
# -----------------------------
|
|
def build_prompt(
|
|
mode: str,
|
|
user_prompt: str,
|
|
grounding: bool,
|
|
find_term: Optional[str],
|
|
schema: Optional[str],
|
|
include_caption: bool,
|
|
) -> str:
|
|
"""Build the prompt based on mode"""
|
|
parts: List[str] = ["<image>"]
|
|
mode_requires_grounding = mode in {"find_ref", "layout_map", "pii_redact"}
|
|
if grounding or mode_requires_grounding:
|
|
parts.append("<|grounding|>")
|
|
|
|
instruction = ""
|
|
if mode == "plain_ocr":
|
|
instruction = "Free OCR."
|
|
elif mode == "markdown":
|
|
instruction = "Convert the document to markdown."
|
|
elif mode == "tables_csv":
|
|
instruction = (
|
|
"Extract every table and output CSV only. "
|
|
"Use commas, minimal quoting. If multiple tables, separate with a line containing '---'."
|
|
)
|
|
elif mode == "tables_md":
|
|
instruction = "Extract every table as GitHub-flavored Markdown tables. Output only the tables."
|
|
elif mode == "kv_json":
|
|
schema_text = schema.strip() if schema else "{}"
|
|
instruction = (
|
|
"Extract key fields and return strict JSON only. "
|
|
f"Use this schema (fill the values): {schema_text}"
|
|
)
|
|
elif mode == "figure_chart":
|
|
instruction = (
|
|
"Parse the figure. First extract any numeric series as a two-column table (x,y). "
|
|
"Then summarize the chart in 2 sentences. Output the table, then a line '---', then the summary."
|
|
)
|
|
elif mode == "find_ref":
|
|
key = (find_term or "").strip() or "Total"
|
|
instruction = f"Locate <|ref|>{key}<|/ref|> in the image."
|
|
elif mode == "layout_map":
|
|
instruction = (
|
|
'Return a JSON array of blocks with fields {"type":["title","paragraph","table","figure"],'
|
|
'"box":[x1,y1,x2,y2]}. Do not include any text content.'
|
|
)
|
|
elif mode == "pii_redact":
|
|
instruction = (
|
|
'Find all occurrences of emails, phone numbers, postal addresses, and IBANs. '
|
|
'Return a JSON array of objects {label, text, box:[x1,y1,x2,y2]}.'
|
|
)
|
|
elif mode == "multilingual":
|
|
instruction = "Free OCR. Detect the language automatically and output in the same script."
|
|
elif mode == "describe":
|
|
instruction = "Describe this image. Focus on visible key elements."
|
|
elif mode == "freeform":
|
|
instruction = user_prompt.strip() if user_prompt else "OCR this image."
|
|
else:
|
|
instruction = "OCR this image."
|
|
|
|
if include_caption and mode not in {"describe"}:
|
|
instruction = instruction + "\nThen add a one-paragraph description of the image."
|
|
|
|
parts.append(instruction)
|
|
return "\n".join(parts)
|
|
|
|
# -----------------------------
|
|
# Grounding parser
|
|
# -----------------------------
|
|
# Match a full detection block and capture the coordinates as the entire list expression
|
|
# Examples of captured coords (including outer brackets):
|
|
# - [[312, 339, 480, 681]]
|
|
# - [[504, 700, 625, 910], [771, 570, 996, 996]]
|
|
# - [[110, 310, 255, 800], [312, 343, 479, 680], ...]
|
|
# Using a greedy bracket capture ensures we include all inner lists up to the last ']' before </|det|>
|
|
DET_BLOCK = re.compile(
|
|
r"<\|ref\|>(?P<label>.*?)<\|/ref\|>\s*<\|det\|>\s*(?P<coords>\[.*\])\s*<\|/det\|>",
|
|
re.DOTALL,
|
|
)
|
|
|
|
def clean_grounding_text(text: str) -> str:
|
|
"""Remove grounding tags from text for display, keeping labels"""
|
|
# Replace <|ref|>label<|/ref|><|det|>[...any nested lists...]<|/det|> with just the label
|
|
cleaned = re.sub(
|
|
r"<\|ref\|>(.*?)<\|/ref\|>\s*<\|det\|>\s*\[.*\]\s*<\|/det\|>",
|
|
r"\1",
|
|
text,
|
|
flags=re.DOTALL,
|
|
)
|
|
# Also remove any standalone grounding tags
|
|
cleaned = re.sub(r"<\|grounding\|>", "", cleaned)
|
|
return cleaned.strip()
|
|
|
|
def parse_detections(text: str, image_width: int, image_height: int) -> List[Dict[str, Any]]:
|
|
"""Parse grounding boxes from text and scale from 0-999 normalized coords to actual image dimensions
|
|
|
|
Handles both single and multiple bounding boxes:
|
|
- Single: <|ref|>label<|/ref|><|det|>[[x1,y1,x2,y2]]<|/det|>
|
|
- Multiple: <|ref|>label<|/ref|><|det|>[[x1,y1,x2,y2], [x1,y1,x2,y2], ...]<|/det|>
|
|
"""
|
|
boxes: List[Dict[str, Any]] = []
|
|
for m in DET_BLOCK.finditer(text or ""):
|
|
label = m.group("label").strip()
|
|
coords_str = m.group("coords").strip()
|
|
|
|
print(f"🔍 DEBUG: Found detection for '{label}'")
|
|
print(f"📦 Raw coords string (with brackets): {coords_str}")
|
|
|
|
try:
|
|
import ast
|
|
|
|
# Parse the full bracket expression directly (handles single and multiple)
|
|
parsed = ast.literal_eval(coords_str)
|
|
|
|
# Normalize to a list of lists
|
|
if (
|
|
isinstance(parsed, list)
|
|
and len(parsed) == 4
|
|
and all(isinstance(n, (int, float)) for n in parsed)
|
|
):
|
|
# Single box provided as [x1,y1,x2,y2]
|
|
box_coords = [parsed]
|
|
print("📦 Single box (flat list) detected")
|
|
elif isinstance(parsed, list):
|
|
box_coords = parsed
|
|
print(f"📦 Boxes detected: {len(box_coords)}")
|
|
else:
|
|
raise ValueError("Unsupported coords structure")
|
|
|
|
# Process each box
|
|
for idx, box in enumerate(box_coords):
|
|
if isinstance(box, (list, tuple)) and len(box) >= 4:
|
|
x1 = int(float(box[0]) / 999 * image_width)
|
|
y1 = int(float(box[1]) / 999 * image_height)
|
|
x2 = int(float(box[2]) / 999 * image_width)
|
|
y2 = int(float(box[3]) / 999 * image_height)
|
|
print(f" Box {idx+1}: {box} → [{x1}, {y1}, {x2}, {y2}]")
|
|
boxes.append({"label": label, "box": [x1, y1, x2, y2]})
|
|
else:
|
|
print(f" ⚠️ Skipping invalid box: {box}")
|
|
except Exception as e:
|
|
print(f"❌ Parsing failed: {e}")
|
|
continue
|
|
|
|
print(f"🎯 Total boxes parsed: {len(boxes)}")
|
|
return boxes
|
|
|
|
# -----------------------------
|
|
# Routes
|
|
# -----------------------------
|
|
@app.get("/")
|
|
async def root():
|
|
return {"message": "DeepSeek-OCR API is running! 🚀", "docs": "/docs"}
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy", "model_loaded": model is not None}
|
|
|
|
@app.post("/api/ocr")
|
|
async def ocr_inference(
|
|
image: UploadFile = File(...),
|
|
mode: str = Form("plain_ocr"),
|
|
prompt: str = Form(""),
|
|
grounding: bool = Form(False),
|
|
include_caption: bool = Form(False),
|
|
find_term: Optional[str] = Form(None),
|
|
schema: Optional[str] = Form(None),
|
|
base_size: int = Form(1024),
|
|
image_size: int = Form(640),
|
|
crop_mode: bool = Form(True),
|
|
test_compress: bool = Form(False),
|
|
):
|
|
"""
|
|
Perform OCR inference on uploaded image
|
|
|
|
- **image**: Image file to process
|
|
- **mode**: OCR mode (plain_ocr, markdown, tables_csv, etc.)
|
|
- **prompt**: Custom prompt for freeform mode
|
|
- **grounding**: Enable grounding boxes
|
|
- **include_caption**: Add image description
|
|
- **find_term**: Term to find (for find_ref mode)
|
|
- **schema**: JSON schema (for kv_json mode)
|
|
- **base_size**: Base processing size
|
|
- **image_size**: Image size parameter
|
|
- **crop_mode**: Enable crop mode
|
|
- **test_compress**: Test compression
|
|
"""
|
|
if model is None or tokenizer is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
# Build prompt
|
|
prompt_text = build_prompt(
|
|
mode=mode,
|
|
user_prompt=prompt,
|
|
grounding=grounding,
|
|
find_term=find_term,
|
|
schema=schema,
|
|
include_caption=include_caption,
|
|
)
|
|
|
|
tmp_img = None
|
|
out_dir = None
|
|
try:
|
|
# Save uploaded file
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
|
|
content = await image.read()
|
|
tmp.write(content)
|
|
tmp_img = tmp.name
|
|
|
|
# Get original dimensions
|
|
try:
|
|
with Image.open(tmp_img) as im:
|
|
orig_w, orig_h = im.size
|
|
except Exception:
|
|
orig_w = orig_h = None
|
|
|
|
out_dir = tempfile.mkdtemp(prefix="dsocr_")
|
|
|
|
# Run inference
|
|
res = model.infer(
|
|
tokenizer,
|
|
prompt=prompt_text,
|
|
image_file=tmp_img,
|
|
output_path=out_dir,
|
|
base_size=base_size,
|
|
image_size=image_size,
|
|
crop_mode=crop_mode,
|
|
save_results=False,
|
|
test_compress=test_compress,
|
|
eval_mode=True,
|
|
)
|
|
|
|
# Normalize response
|
|
if isinstance(res, str):
|
|
text = res.strip()
|
|
elif isinstance(res, dict) and "text" in res:
|
|
text = str(res["text"]).strip()
|
|
elif isinstance(res, (list, tuple)):
|
|
text = "\n".join(map(str, res)).strip()
|
|
else:
|
|
text = ""
|
|
|
|
# Fallback: check output file
|
|
if not text:
|
|
mmd = os.path.join(out_dir, "result.mmd")
|
|
if os.path.exists(mmd):
|
|
with open(mmd, "r", encoding="utf-8") as fh:
|
|
text = fh.read().strip()
|
|
if not text:
|
|
text = "No text returned by model."
|
|
|
|
# Parse grounding boxes with proper coordinate scaling
|
|
boxes = parse_detections(text, orig_w or 1, orig_h or 1) if ("<|det|>" in text or "<|ref|>" in text) else []
|
|
|
|
# Clean grounding tags from display text, but keep the labels
|
|
display_text = clean_grounding_text(text) if ("<|ref|>" in text or "<|grounding|>" in text) else text
|
|
|
|
# If display text is empty after cleaning but we have boxes, show the labels
|
|
if not display_text and boxes:
|
|
display_text = ", ".join([b["label"] for b in boxes])
|
|
|
|
return JSONResponse({
|
|
"success": True,
|
|
"text": display_text,
|
|
"raw_text": text, # Include raw model output for debugging
|
|
"boxes": boxes,
|
|
"image_dims": {"w": orig_w, "h": orig_h},
|
|
"metadata": {
|
|
"mode": mode,
|
|
"grounding": grounding or (mode in {"find_ref","layout_map","pii_redact"}),
|
|
"base_size": base_size,
|
|
"image_size": image_size,
|
|
"crop_mode": crop_mode
|
|
}
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"OCR inference error: {type(e).__name__}: {str(e)}")
|
|
raise HTTPException(status_code=500, detail="An internal error occurred during OCR processing.")
|
|
|
|
finally:
|
|
if tmp_img:
|
|
try:
|
|
os.remove(tmp_img)
|
|
except Exception:
|
|
pass
|
|
if out_dir:
|
|
shutil.rmtree(out_dir, ignore_errors=True)
|
|
|
|
@app.post("/api/process-pdf")
|
|
async def process_pdf(
|
|
pdf_file: UploadFile = File(...),
|
|
mode: str = Form("plain_ocr"),
|
|
prompt: str = Form(""),
|
|
output_format: str = Form("markdown"), # markdown, html, docx, json
|
|
grounding: bool = Form(False),
|
|
include_caption: bool = Form(False),
|
|
extract_images: bool = Form(True),
|
|
dpi: int = Form(144),
|
|
base_size: int = Form(1024),
|
|
image_size: int = Form(640),
|
|
crop_mode: bool = Form(True),
|
|
):
|
|
"""
|
|
Process PDF document with OCR and convert to various formats
|
|
|
|
- **pdf_file**: PDF file to process
|
|
- **mode**: OCR mode (plain_ocr, markdown, tables_csv, etc.)
|
|
- **prompt**: Custom prompt for freeform mode
|
|
- **output_format**: Output format (markdown, html, docx, json)
|
|
- **grounding**: Enable grounding boxes
|
|
- **include_caption**: Add image descriptions
|
|
- **extract_images**: Extract images from PDF
|
|
- **dpi**: PDF rendering resolution (default: 144)
|
|
- **base_size**: Base processing size
|
|
- **image_size**: Image size parameter
|
|
- **crop_mode**: Enable crop mode
|
|
"""
|
|
if model is None or tokenizer is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
# Validate output format
|
|
if output_format not in ["markdown", "html", "docx", "json"]:
|
|
raise HTTPException(status_code=400, detail="Invalid output format. Must be: markdown, html, docx, or json")
|
|
|
|
try:
|
|
# Read PDF file
|
|
pdf_bytes = await pdf_file.read()
|
|
|
|
# Convert PDF to images
|
|
print(f"📄 Converting PDF to images (DPI: {dpi})...")
|
|
images = pdf_to_images_high_quality(pdf_bytes, dpi=dpi)
|
|
total_pages = len(images)
|
|
print(f"✅ Converted {total_pages} pages")
|
|
|
|
# Process each page
|
|
pages_content = []
|
|
converter = DocumentConverter()
|
|
|
|
for page_idx, img in enumerate(images):
|
|
print(f"🔍 Processing page {page_idx + 1}/{total_pages}...")
|
|
|
|
# Build prompt for this page
|
|
prompt_text = build_prompt(
|
|
mode=mode,
|
|
user_prompt=prompt,
|
|
grounding=grounding,
|
|
find_term=None,
|
|
schema=None,
|
|
include_caption=include_caption,
|
|
)
|
|
|
|
# Save image temporarily
|
|
tmp_img = None
|
|
out_dir = None
|
|
try:
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
|
|
img.save(tmp, format="PNG")
|
|
tmp_img = tmp.name
|
|
|
|
orig_w, orig_h = img.size
|
|
out_dir = tempfile.mkdtemp(prefix="dsocr_pdf_")
|
|
|
|
# Run inference
|
|
res = model.infer(
|
|
tokenizer,
|
|
prompt=prompt_text,
|
|
image_file=tmp_img,
|
|
output_path=out_dir,
|
|
base_size=base_size,
|
|
image_size=image_size,
|
|
crop_mode=crop_mode,
|
|
save_results=False,
|
|
test_compress=False,
|
|
eval_mode=True,
|
|
)
|
|
|
|
# Normalize response
|
|
if isinstance(res, str):
|
|
text = res.strip()
|
|
elif isinstance(res, dict) and "text" in res:
|
|
text = str(res["text"]).strip()
|
|
elif isinstance(res, (list, tuple)):
|
|
text = "\n".join(map(str, res)).strip()
|
|
else:
|
|
text = ""
|
|
|
|
if not text:
|
|
mmd = os.path.join(out_dir, "result.mmd")
|
|
if os.path.exists(mmd):
|
|
with open(mmd, "r", encoding="utf-8") as fh:
|
|
text = fh.read().strip()
|
|
if not text:
|
|
text = f"No text returned for page {page_idx + 1}."
|
|
|
|
# Extract images if requested
|
|
page_images = []
|
|
if extract_images:
|
|
matches, matches_image, matches_other = extract_ref_patterns(text)
|
|
if matches_image:
|
|
cropped = crop_images_from_refs(img, matches)
|
|
for cropped_img in cropped:
|
|
# Convert to base64
|
|
img_buffer = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
|
|
cropped_img.save(img_buffer.name, format="JPEG", quality=95)
|
|
with open(img_buffer.name, "rb") as f:
|
|
img_b64 = base64.b64encode(f.read()).decode('utf-8')
|
|
page_images.append(img_b64)
|
|
os.remove(img_buffer.name)
|
|
|
|
# Clean the text and add image placeholders
|
|
text = clean_markdown_content(text, matches_image, matches_other)
|
|
for img_idx in range(len(page_images)):
|
|
text = f"[IMAGE_{img_idx}]\n" + text
|
|
|
|
# Parse grounding boxes
|
|
boxes = parse_detections(text, orig_w, orig_h) if ("<|det|>" in text or "<|ref|>" in text) else []
|
|
|
|
# Clean grounding tags from display text
|
|
display_text = clean_grounding_text(text) if ("<|ref|>" in text or "<|grounding|>" in text) else text
|
|
|
|
pages_content.append({
|
|
'page_number': page_idx + 1,
|
|
'text': display_text,
|
|
'raw_text': text,
|
|
'boxes': boxes,
|
|
'images': page_images,
|
|
'image_dims': {'w': orig_w, 'h': orig_h}
|
|
})
|
|
|
|
finally:
|
|
if tmp_img:
|
|
try:
|
|
os.remove(tmp_img)
|
|
except Exception:
|
|
pass
|
|
if out_dir:
|
|
shutil.rmtree(out_dir, ignore_errors=True)
|
|
|
|
print(f"✅ Processed all {total_pages} pages")
|
|
|
|
# Convert to requested format
|
|
if output_format == "json":
|
|
return JSONResponse({
|
|
"success": True,
|
|
"total_pages": total_pages,
|
|
"pages": pages_content,
|
|
"metadata": {
|
|
"mode": mode,
|
|
"grounding": grounding,
|
|
"extract_images": extract_images,
|
|
"dpi": dpi
|
|
}
|
|
})
|
|
elif output_format == "markdown":
|
|
md_content = converter.to_markdown(pages_content, include_images=extract_images)
|
|
return StreamingResponse(
|
|
iter([md_content.encode('utf-8')]),
|
|
media_type="text/markdown",
|
|
headers={"Content-Disposition": f"attachment; filename=ocr_result.md"}
|
|
)
|
|
elif output_format == "html":
|
|
html_content = converter.to_html(pages_content, include_images=extract_images)
|
|
return StreamingResponse(
|
|
iter([html_content.encode('utf-8')]),
|
|
media_type="text/html",
|
|
headers={"Content-Disposition": f"attachment; filename=ocr_result.html"}
|
|
)
|
|
elif output_format == "docx":
|
|
docx_buffer = converter.to_docx(pages_content, include_images=extract_images)
|
|
return StreamingResponse(
|
|
docx_buffer,
|
|
media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
headers={"Content-Disposition": f"attachment; filename=ocr_result.docx"}
|
|
)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
print(f"Error processing PDF: {e}")
|
|
print(traceback.format_exc())
|
|
raise HTTPException(status_code=500, detail="An internal error occurred during PDF processing.")
|
|
|
|
# -----------------------------
|
|
# Job management routes
|
|
# -----------------------------
|
|
|
|
class ReviewRequest(BaseModel):
|
|
reviewed_text: str
|
|
reviewer_name: str
|
|
author: Optional[str] = None
|
|
book: Optional[str] = None
|
|
chapter: Optional[str] = None
|
|
page: Optional[str] = None
|
|
|
|
|
|
def _job_row_to_dict(row) -> Dict[str, Any]:
|
|
"""Convert a DB row (RealDictRow) to a plain dict with serialisable values."""
|
|
d = dict(row)
|
|
for key, val in d.items():
|
|
if isinstance(val, datetime):
|
|
d[key] = val.isoformat()
|
|
elif val is not None and hasattr(val, '__str__') and type(val).__name__ == 'UUID':
|
|
d[key] = str(val)
|
|
return d
|
|
|
|
|
|
@app.post("/api/jobs")
|
|
async def commit_job(
|
|
image: UploadFile = File(...),
|
|
author: str = Form(""),
|
|
book: str = Form(""),
|
|
chapter: str = Form(""),
|
|
page: str = Form(""),
|
|
ocr_text: str = Form(""),
|
|
mode: str = Form("plain_ocr"),
|
|
):
|
|
"""Commit an OCR job: save the image and insert a DB record."""
|
|
job_id = str(uuid.uuid4())
|
|
|
|
# Determine file extension from original filename or content type
|
|
original_filename = image.filename or "image"
|
|
ext = os.path.splitext(original_filename)[1].lower()
|
|
if not ext:
|
|
ct = (image.content_type or "").lower()
|
|
ext_map = {
|
|
"image/png": ".png", "image/jpeg": ".jpg", "image/jpg": ".jpg",
|
|
"image/webp": ".webp", "image/gif": ".gif", "image/bmp": ".bmp",
|
|
}
|
|
ext = ext_map.get(ct, ".png")
|
|
|
|
image_path = os.path.join(OCR_IMAGES_DIR, f"{job_id}{ext}")
|
|
|
|
try:
|
|
content = await image.read()
|
|
with open(image_path, "wb") as f:
|
|
f.write(content)
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=500, detail="Failed to save image file.")
|
|
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO ocr_jobs
|
|
(id, author, book, chapter, page, image_path, original_filename,
|
|
ocr_text, mode, status)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, 'unreviewed')
|
|
RETURNING *
|
|
""",
|
|
(job_id, author or None, book or None, chapter or None,
|
|
page or None, image_path, original_filename,
|
|
ocr_text or None, mode),
|
|
)
|
|
row = cur.fetchone()
|
|
except Exception as exc:
|
|
# Clean up saved image if DB insert fails
|
|
try:
|
|
os.remove(image_path)
|
|
except Exception:
|
|
pass
|
|
# Unique constraint violation (author + chapter + page already exists)
|
|
if getattr(exc, 'pgcode', None) == '23505':
|
|
raise HTTPException(
|
|
status_code=409,
|
|
detail="A job with this Author, Chapter, and Page already exists."
|
|
)
|
|
print(f"Job commit DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Failed to save job to database.")
|
|
|
|
return JSONResponse(_job_row_to_dict(row), status_code=201)
|
|
|
|
|
|
@app.get("/api/jobs")
|
|
async def list_jobs(
|
|
search: Optional[str] = Query(None, description="General text search across all fields"),
|
|
author: Optional[str] = Query(None),
|
|
book: Optional[str] = Query(None),
|
|
chapter: Optional[str] = Query(None),
|
|
status: Optional[str] = Query(None, description="unreviewed | reviewed"),
|
|
limit: int = Query(20, ge=1, le=200),
|
|
offset: int = Query(0, ge=0),
|
|
):
|
|
"""Search and list jobs. All filters are optional and combinable."""
|
|
conditions = []
|
|
params: List[Any] = []
|
|
|
|
if search:
|
|
conditions.append(
|
|
"(author ILIKE %s OR book ILIKE %s OR chapter ILIKE %s "
|
|
"OR page ILIKE %s OR ocr_text ILIKE %s OR reviewer_name ILIKE %s)"
|
|
)
|
|
like = f"%{search}%"
|
|
params.extend([like, like, like, like, like, like])
|
|
|
|
if author:
|
|
conditions.append("author ILIKE %s")
|
|
params.append(f"%{author}%")
|
|
|
|
if book:
|
|
conditions.append("book ILIKE %s")
|
|
params.append(f"%{book}%")
|
|
|
|
if chapter:
|
|
conditions.append("chapter ILIKE %s")
|
|
params.append(f"%{chapter}%")
|
|
|
|
if status:
|
|
conditions.append("status = %s")
|
|
params.append(status)
|
|
|
|
where = ("WHERE " + " AND ".join(conditions)) if conditions else ""
|
|
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
f"SELECT COUNT(*) AS total FROM ocr_jobs {where}",
|
|
params,
|
|
)
|
|
total = cur.fetchone()["total"]
|
|
|
|
cur.execute(
|
|
f"""
|
|
SELECT id, author, book, chapter, page, submitted_at, status,
|
|
reviewer_name, reviewed_at, mode, original_filename
|
|
FROM ocr_jobs {where}
|
|
ORDER BY submitted_at DESC
|
|
LIMIT %s OFFSET %s
|
|
""",
|
|
params + [limit, offset],
|
|
)
|
|
rows = [_job_row_to_dict(r) for r in cur.fetchall()]
|
|
except Exception as exc:
|
|
print(f"list_jobs DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Database error.")
|
|
|
|
return JSONResponse({"total": total, "limit": limit, "offset": offset, "jobs": rows})
|
|
|
|
|
|
@app.get("/api/jobs/suggestions")
|
|
async def job_suggestions():
|
|
"""Return distinct values for author, chapter, and reviewer_name to power autocomplete."""
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("""
|
|
SELECT
|
|
array_remove(array_agg(DISTINCT author ORDER BY author), NULL) AS authors,
|
|
array_remove(array_agg(DISTINCT chapter ORDER BY chapter), NULL) AS chapters,
|
|
array_remove(array_agg(DISTINCT reviewer_name ORDER BY reviewer_name), NULL) AS reviewers
|
|
FROM ocr_jobs
|
|
""")
|
|
row = cur.fetchone()
|
|
except Exception as exc:
|
|
print(f"suggestions DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Database error.")
|
|
|
|
return JSONResponse({
|
|
"authors": row["authors"] or [],
|
|
"chapters": row["chapters"] or [],
|
|
"reviewers": row["reviewers"] or [],
|
|
})
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}")
|
|
async def get_job(job_id: str):
|
|
"""Retrieve full job record including OCR text."""
|
|
try:
|
|
uuid.UUID(job_id)
|
|
except ValueError:
|
|
raise HTTPException(status_code=400, detail="Invalid job ID.")
|
|
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT * FROM ocr_jobs WHERE id = %s", (job_id,))
|
|
row = cur.fetchone()
|
|
except Exception as exc:
|
|
print(f"get_job DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Database error.")
|
|
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail="Job not found.")
|
|
|
|
return JSONResponse(_job_row_to_dict(row))
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}/image")
|
|
async def get_job_image(job_id: str):
|
|
"""Serve the stored image for a job."""
|
|
try:
|
|
uuid.UUID(job_id)
|
|
except ValueError:
|
|
raise HTTPException(status_code=400, detail="Invalid job ID.")
|
|
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT image_path FROM ocr_jobs WHERE id = %s", (job_id,))
|
|
row = cur.fetchone()
|
|
except Exception as exc:
|
|
print(f"get_job_image DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Database error.")
|
|
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail="Job not found.")
|
|
|
|
path = row["image_path"]
|
|
if not os.path.isfile(path):
|
|
raise HTTPException(status_code=404, detail="Image file not found on disk.")
|
|
|
|
return FileResponse(path)
|
|
|
|
|
|
@app.put("/api/jobs/{job_id}/review")
|
|
async def review_job(job_id: str, body: ReviewRequest):
|
|
"""Mark a job as reviewed with the corrected text and reviewer name."""
|
|
try:
|
|
uuid.UUID(job_id)
|
|
except ValueError:
|
|
raise HTTPException(status_code=400, detail="Invalid job ID.")
|
|
|
|
try:
|
|
with get_db() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE ocr_jobs
|
|
SET status = 'reviewed',
|
|
reviewed_text = %s,
|
|
reviewer_name = %s,
|
|
reviewed_at = NOW(),
|
|
author = %s,
|
|
book = %s,
|
|
chapter = %s,
|
|
page = %s
|
|
WHERE id = %s
|
|
RETURNING *
|
|
""",
|
|
(
|
|
body.reviewed_text,
|
|
body.reviewer_name,
|
|
body.author or None,
|
|
body.book or None,
|
|
body.chapter or None,
|
|
body.page or None,
|
|
job_id,
|
|
),
|
|
)
|
|
row = cur.fetchone()
|
|
except Exception as exc:
|
|
print(f"review_job DB error: {exc}")
|
|
raise HTTPException(status_code=500, detail="Database error.")
|
|
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail="Job not found.")
|
|
|
|
return JSONResponse(_job_row_to_dict(row))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
host = env_config("API_HOST", default="0.0.0.0")
|
|
port = env_config("API_PORT", default=8000, cast=int)
|
|
uvicorn.run(app, host=host, port=port)
|