Spaces:
Paused
Paused
File size: 3,215 Bytes
9251dad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from typing import Optional, Union
import base64
from .models import OCRRequest, SlideMatchRequest, DetectionRequest, APIResponse
from .services import ocr_service
app = FastAPI()
from starlette.datastructures import UploadFile as StarletteUploadFile
async def decode_image(image: Union[UploadFile, StarletteUploadFile, str, None]) -> bytes:
if image is None:
raise HTTPException(status_code=400, detail="No image provided")
if isinstance(image, (UploadFile, StarletteUploadFile)):
return await image.read()
elif isinstance(image, str):
try:
# 检查是否是 base64 编码的图片
if image.startswith(('data:image/', 'data:application/')):
# 移除 MIME 类型前缀
image = image.split(',')[1]
return base64.b64decode(image)
except:
raise HTTPException(status_code=400, detail="Invalid base64 string")
else:
raise HTTPException(status_code=400, detail="Invalid image input")
@app.post("/ocr", response_model=APIResponse)
async def ocr_endpoint(
file: Optional[UploadFile] = File(None),
image: Optional[str] = Form(None),
probability: bool = Form(False),
charsets: Optional[str] = Form(None),
png_fix: bool = Form(False)
):
try:
if file is None and image is None:
return APIResponse(code=400, message="Either file or image must be provided")
image_bytes = await decode_image(file or image)
result = ocr_service.ocr_classification(image_bytes, probability, charsets, png_fix)
return APIResponse(code=200, message="Success", data=result)
except Exception as e:
return APIResponse(code=500, message=str(e))
@app.post("/slide_match", response_model=APIResponse)
async def slide_match_endpoint(
target_file: Optional[UploadFile] = File(None),
background_file: Optional[UploadFile] = File(None),
target: Optional[str] = Form(None),
background: Optional[str] = Form(None),
simple_target: bool = Form(False)
):
try:
target_bytes = await decode_image(target_file or target)
background_bytes = await decode_image(background_file or background)
result = ocr_service.slide_match(target_bytes, background_bytes, simple_target)
return APIResponse(code=200, message="Success", data=result)
except Exception as e:
return APIResponse(code=500, message=str(e))
@app.post("/detection", response_model=APIResponse)
async def detection_endpoint(
file: Optional[UploadFile] = File(None),
image: Optional[str] = Form(None)
):
try:
if file is None and image is None:
return APIResponse(code=400, message="Either file or image must be provided")
image_bytes = await decode_image(file or image)
bboxes = ocr_service.detection(image_bytes)
return APIResponse(code=200, message="Success", data=bboxes)
except Exception as e:
return APIResponse(code=500, message=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000) |