mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 06:05:45 +08:00
Refactor: migrate pdf_parser.py to golang (#16323)
### What problem does this PR solve? Http API based on onnx model. pdf_parser.py to golang ### Type of change - [x] Refactoring
This commit is contained in:
0
deepdoc/server/endpoints/__init__.py
Normal file
0
deepdoc/server/endpoints/__init__.py
Normal file
43
deepdoc/server/endpoints/dla_endpoint.py
Normal file
43
deepdoc/server/endpoints/dla_endpoint.py
Normal file
@@ -0,0 +1,43 @@
|
||||
"""DLA LitServe endpoint."""
|
||||
|
||||
import logging
|
||||
|
||||
import litserve as ls
|
||||
|
||||
from deepdoc.server.adapters.dla_adapter import DLAAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DLAEndpoint(ls.LitAPI):
|
||||
"""Document Layout Analysis endpoint at /predict/dla."""
|
||||
|
||||
def __init__(self, model_dir: str, thr: float = 0.2):
|
||||
super().__init__()
|
||||
self.api_path = "/predict/dla"
|
||||
self.model_dir = model_dir
|
||||
self.thr = thr
|
||||
self.adapter: DLAAdapter | None = None
|
||||
|
||||
def setup(self, device):
|
||||
self.adapter = DLAAdapter(model_dir=self.model_dir, thr=self.thr)
|
||||
self.adapter.load()
|
||||
logger.info("DLA model loaded")
|
||||
|
||||
def decode_request(self, request):
|
||||
# Handle both Starlette UploadFile (old) and FormData (Starlette >=1.3)
|
||||
if hasattr(request, "file"):
|
||||
data = request.file.read()
|
||||
else:
|
||||
data = request.get("request").file.read()
|
||||
if not data:
|
||||
raise ValueError("Empty request body")
|
||||
if len(data) > 50 * 1024 * 1024: # 50MB
|
||||
raise ValueError("Image too large")
|
||||
return data
|
||||
|
||||
def predict(self, image_data: bytes):
|
||||
return self.adapter(image_data)
|
||||
|
||||
def encode_response(self, output):
|
||||
return {"bboxes": output}
|
||||
67
deepdoc/server/endpoints/ocr_endpoint.py
Normal file
67
deepdoc/server/endpoints/ocr_endpoint.py
Normal file
@@ -0,0 +1,67 @@
|
||||
"""OCR LitServe endpoint — detect + rec via operator form field."""
|
||||
|
||||
import logging
|
||||
|
||||
import litserve as ls
|
||||
|
||||
from deepdoc.server.adapters.ocr_adapter import OCRAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OCREndpoint(ls.LitAPI):
|
||||
"""OCR endpoint at /predict/ocr.
|
||||
|
||||
Form field 'operator' (det or rec) selects the mode.
|
||||
Form field 'request' carries the JPEG image bytes.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir: str):
|
||||
super().__init__()
|
||||
self.api_path = "/predict/ocr"
|
||||
self.model_dir = model_dir
|
||||
self.adapter: OCRAdapter | None = None
|
||||
|
||||
def setup(self, device):
|
||||
self.adapter = OCRAdapter(model_dir=self.model_dir)
|
||||
self.adapter.load()
|
||||
logger.info("OCR model loaded")
|
||||
|
||||
def decode_request(self, request):
|
||||
# Handle both old Starlette UploadFile and new Starlette FormData
|
||||
if hasattr(request, "file"):
|
||||
data = request.file.read()
|
||||
# Try to read operator from the underlying request context
|
||||
operator = getattr(self, "_request", None)
|
||||
if operator is not None:
|
||||
operator = operator.query_params.get("operator", "")
|
||||
else:
|
||||
operator = ""
|
||||
else:
|
||||
# FormData: get file and operator form fields
|
||||
data = request.get("request").file.read()
|
||||
op_val = request.get("operator")
|
||||
operator = str(op_val) if op_val else ""
|
||||
|
||||
if not data:
|
||||
raise ValueError("Empty request body")
|
||||
if len(data) > 50 * 1024 * 1024:
|
||||
raise ValueError("Image too large")
|
||||
|
||||
operator = operator.strip().lower()
|
||||
if operator not in ("det", "rec"):
|
||||
raise ValueError(
|
||||
f"Invalid or missing operator '{operator}' (must be 'det' or 'rec')"
|
||||
)
|
||||
|
||||
return operator, data
|
||||
|
||||
def predict(self, inputs: tuple):
|
||||
operator, image_data = inputs
|
||||
if operator == "det":
|
||||
return self.adapter.detect(image_data)
|
||||
else:
|
||||
return self.adapter.recognize(image_data)
|
||||
|
||||
def encode_response(self, output):
|
||||
return output
|
||||
43
deepdoc/server/endpoints/tsr_endpoint.py
Normal file
43
deepdoc/server/endpoints/tsr_endpoint.py
Normal file
@@ -0,0 +1,43 @@
|
||||
"""TSR LitServe endpoint."""
|
||||
|
||||
import logging
|
||||
|
||||
import litserve as ls
|
||||
|
||||
from deepdoc.server.adapters.tsr_adapter import TSRAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TSREndpoint(ls.LitAPI):
|
||||
"""Table Structure Recognition endpoint at /predict/tsr."""
|
||||
|
||||
def __init__(self, model_dir: str, thr: float = 0.2):
|
||||
super().__init__()
|
||||
self.api_path = "/predict/tsr"
|
||||
self.model_dir = model_dir
|
||||
self.thr = thr
|
||||
self.adapter: TSRAdapter | None = None
|
||||
|
||||
def setup(self, device):
|
||||
self.adapter = TSRAdapter(model_dir=self.model_dir, thr=self.thr)
|
||||
self.adapter.load()
|
||||
logger.info("TSR model loaded")
|
||||
|
||||
def decode_request(self, request):
|
||||
# Handle both Starlette UploadFile (old) and FormData (Starlette >=1.3)
|
||||
if hasattr(request, "file"):
|
||||
data = request.file.read()
|
||||
else:
|
||||
data = request.get("request").file.read()
|
||||
if not data:
|
||||
raise ValueError("Empty request body")
|
||||
if len(data) > 50 * 1024 * 1024:
|
||||
raise ValueError("Image too large")
|
||||
return data
|
||||
|
||||
def predict(self, image_data: bytes):
|
||||
return self.adapter(image_data)
|
||||
|
||||
def encode_response(self, output):
|
||||
return {"bboxes": output}
|
||||
Reference in New Issue
Block a user