File size: 2,766 Bytes
c8cdee4 | 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 | """Local browser UI. Uploaded images live only in request memory."""
import argparse
import io
import threading
import time
import webbrowser
from pathlib import Path
from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.responses import FileResponse
from PIL import Image, UnidentifiedImageError
from starlette.concurrency import run_in_threadpool
from runtime.engine import OCR
def create_app(ocr):
app=FastAPI(title='Bashkir OCR Studio',docs_url=None,redoc_url=None)
lock=threading.Lock()
@app.get('/')
def index():return FileResponse(Path(__file__).parent/'runtime/index.html',headers={'Cache-Control':'no-store, no-cache, must-revalidate','Pragma':'no-cache','Expires':'0'})
@app.get('/status')
def status():return dict(ready=True,backend=ocr.backend)
@app.post('/recognize')
async def recognize(file:UploadFile):
data=await file.read(25*1024*1024+1)
if len(data)>25*1024*1024:raise HTTPException(413,'Изображение слишком большое: предел 25 МБ.')
try:
image=Image.open(io.BytesIO(data))
if image.width*image.height>25_000_000:raise HTTPException(413,'Предел — 25 миллионов пикселей.')
image.load()
except (UnidentifiedImageError,OSError,Image.DecompressionBombError):raise HTTPException(400,'Не удалось открыть изображение.')
if not lock.acquire(blocking=False):raise HTTPException(409,'Дождитесь завершения текущего изображения.')
try:return await run_in_threadpool(ocr.recognize,image)
finally:lock.release()
return app
def main():
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--backend',choices=['cpu','gpu','auto'],default='cpu')
parser.add_argument('--port',type=int,default=8001)
parser.add_argument('--no-browser',action='store_true')
parser.add_argument('--model-dir',type=Path,help='Optional directory containing config.json and model files')
args=parser.parse_args()
ocr=OCR(root=args.model_dir,backend=args.backend)
if not args.no_browser:
def open_ready():
import urllib.request
for _ in range(60):
try:
with urllib.request.urlopen(f'http://127.0.0.1:{args.port}/status',timeout=1) as r:
if r.status==200:webbrowser.open(f'http://127.0.0.1:{args.port}/?v={int(time.time())}');return
except OSError:time.sleep(.5)
threading.Thread(target=open_ready,daemon=True).start()
import uvicorn
uvicorn.run(create_app(ocr),host='127.0.0.1',port=args.port,access_log=False)
if __name__=='__main__':main()
|