Fix buffer-size error, fix unwrap and illumination endpoint no work, change .png to .jpg ansew format

This commit is contained in:
2026-09-17 07:11:34 +03:00
parent 155f2f9acf
commit 257850c869

61
api.py
View File

@ -137,7 +137,8 @@ def process_unwrap(net, payload: bytes) -> bytes:
result = np.clip(result, 0, 255).astype(np.uint8) result = np.clip(result, 0, 255).astype(np.uint8)
result_bgr = np.ascontiguousarray(result[:, :, ::-1]) result_bgr = np.ascontiguousarray(result[:, :, ::-1])
ok, buf = cv2.imencode(".png", result_bgr) # ИЗМЕНЕНО: кодирование в JPEG с качеством 100%
ok, buf = cv2.imencode(".jpg", result_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
if not ok: if not ok:
raise HTTPException(status_code=500, detail="Failed to encode result image") raise HTTPException(status_code=500, detail="Failed to encode result image")
return buf.tobytes() return buf.tobytes()
@ -263,7 +264,8 @@ def process_docres(model, payload: bytes, task: str) -> bytes:
except Exception as exc: except Exception as exc:
raise HTTPException(status_code=500, detail=f"{task} failed: {exc}") from exc raise HTTPException(status_code=500, detail=f"{task} failed: {exc}") from exc
ok, buf = cv2.imencode(".png", np.ascontiguousarray(result)) # ИЗМЕНЕНО: кодирование в JPEG с качеством 100%
ok, buf = cv2.imencode(".jpg", np.ascontiguousarray(result), [int(cv2.IMWRITE_JPEG_QUALITY), 100])
if not ok: if not ok:
raise HTTPException(status_code=500, detail="Failed to encode result image") raise HTTPException(status_code=500, detail="Failed to encode result image")
return buf.tobytes() return buf.tobytes()
@ -316,50 +318,51 @@ def health():
} }
} }
@app.post("/unwrap") @app.post("/unwrap")
def unwrap(request: Request, file: UploadFile = File(...)): async def unwrap(request: Request, file: UploadFile = File(...)):
try: # Читаем асинхронно — Starlette сам корректно обработает SpooledTemporaryFile
file.file.seek(0) payload = await file.read()
payload = file.file.read()
except Exception as exc:
raise HTTPException(status_code=400, detail="Cannot read uploaded file") from exc
if not payload: if not payload:
raise HTTPException(status_code=400, detail="Uploaded file is empty") raise HTTPException(status_code=400, detail="Uploaded file is empty")
result_png = process_unwrap(request.app.state.docscanner_net, payload) # Тяжёлую работу отправляем в threadpool, чтобы не блокировать event loop
import asyncio
loop = asyncio.get_running_loop()
result_jpeg = await loop.run_in_executor(
None, process_unwrap, request.app.state.docscanner_net, payload
)
return Response( return Response(
content=result_png, content=result_jpeg,
media_type="image/png", media_type="image/jpeg",
headers={"Content-Disposition": 'inline; filename="unwrap.png"', "Cache-Control": "no-store"}, headers={"Content-Disposition": 'inline; filename="unwrap.jpg"', "Cache-Control": "no-store"},
) )
def _docres_endpoint(request: Request, file: UploadFile, task: str): async def _docres_endpoint_async(request: Request, file: UploadFile, task: str):
try: payload = await file.read()
file.file.seek(0)
payload = file.file.read()
except Exception as exc:
raise HTTPException(status_code=400, detail="Cannot read uploaded file") from exc
if not payload: if not payload:
raise HTTPException(status_code=400, detail="Uploaded file is empty") raise HTTPException(status_code=400, detail="Uploaded file is empty")
result_png = process_docres(request.app.state.docres_model, payload, task) import asyncio
loop = asyncio.get_running_loop()
result_jpeg = await loop.run_in_executor(
None, process_docres, request.app.state.docres_model, payload, task
)
return Response( return Response(
content=result_png, content=result_jpeg,
media_type="image/png", media_type="image/jpeg",
headers={"Content-Disposition": f'inline; filename="{task}.png"', "Cache-Control": "no-store"}, headers={"Content-Disposition": f'inline; filename="{task}.jpg"', "Cache-Control": "no-store"},
) )
@app.post("/deblur") @app.post("/deblur")
def deblur(request: Request, file: UploadFile = File(...)): async def deblur(request: Request, file: UploadFile = File(...)):
return _docres_endpoint(request, file, "deblur") return await _docres_endpoint_async(request, file, "deblur")
@app.post("/illumination_correct") @app.post("/illumination_correct")
def illumination_correct(request: Request, file: UploadFile = File(...)): async def illumination_correct(request: Request, file: UploadFile = File(...)):
return _docres_endpoint(request, file, "illumination_correct") return await _docres_endpoint_async(request, file, "illumination_correct")
if __name__ == "__main__": if __name__ == "__main__":