106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
import io
|
|
from typing import Optional
|
|
from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
from pypdf import PdfReader
|
|
import pytesseract
|
|
from PIL import Image
|
|
|
|
app = FastAPI()
|
|
|
|
# CORS configuration
|
|
origins = [
|
|
"http://localhost",
|
|
"http://localhost:3000",
|
|
"http://localhost:5173", # Vite default
|
|
"http://localhost:4200", # Angular default
|
|
"http://127.0.0.1:4200",
|
|
]
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# Models
|
|
class LoginRequest(BaseModel):
|
|
username: str
|
|
password: str
|
|
|
|
class LoginResponse(BaseModel):
|
|
token: str
|
|
message: str
|
|
|
|
class NERResponse(BaseModel):
|
|
text: str
|
|
|
|
# 1. User Login Module (Dummy)
|
|
@app.post("/api/login", response_model=LoginResponse)
|
|
async def login(request: LoginRequest):
|
|
if request.username == "admin" and request.password == "admin":
|
|
return LoginResponse(token="dummy-jwt-token-123", message="Login Successful")
|
|
raise HTTPException(status_code=401, detail="Invalid credentials")
|
|
|
|
# 2. OCR Module
|
|
def extract_text_from_pdf(file_bytes: bytes) -> str:
|
|
try:
|
|
reader = PdfReader(io.BytesIO(file_bytes))
|
|
text = ""
|
|
for page in reader.pages:
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text += page_text + "\n"
|
|
return text.strip()
|
|
except Exception as e:
|
|
print(f"Error reading PDF: {e}")
|
|
return ""
|
|
|
|
def extract_text_from_image(file_bytes: bytes) -> str:
|
|
try:
|
|
image = Image.open(io.BytesIO(file_bytes))
|
|
text = pytesseract.image_to_string(image)
|
|
return text.strip()
|
|
except Exception as e:
|
|
print(f"Error reading Image: {e}")
|
|
return ""
|
|
|
|
from pdf2image import convert_from_bytes
|
|
|
|
# ... (imports)
|
|
|
|
@app.post("/api/ocr/extract", response_model=NERResponse)
|
|
async def extract_text(file: UploadFile = File(...)):
|
|
content = await file.read()
|
|
filename = file.filename.lower()
|
|
|
|
extracted_text = ""
|
|
|
|
if filename.endswith(".pdf"):
|
|
# Try text extraction first
|
|
extracted_text = extract_text_from_pdf(content)
|
|
|
|
# If text is empty, it might be a scanned PDF.
|
|
if not extracted_text.strip():
|
|
try:
|
|
images = convert_from_bytes(content)
|
|
for i, image in enumerate(images):
|
|
page_text = pytesseract.image_to_string(image)
|
|
extracted_text += f"\n--- Page {i+1} ---\n{page_text}"
|
|
except Exception as e:
|
|
extracted_text = f"Error processing Scanned PDF: {str(e)}\n\n(Hint: Ensure 'poppler' is installed on your system. Run 'brew install poppler')"
|
|
|
|
elif filename.endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp")):
|
|
extracted_text = extract_text_from_image(content)
|
|
else:
|
|
raise HTTPException(status_code=400, detail="Unsupported file type")
|
|
|
|
return NERResponse(text=extracted_text)
|
|
|
|
@app.get("/")
|
|
def read_root():
|
|
return {"message": "OCR Backend API is running"}
|