beerds/server/infra/web/recognizer.py

68 lines
2.4 KiB
Python

import uuid
from dataclasses import dataclass
from typing import Annotated
import inject
from litestar import Controller, Request, Response, get, post
from litestar.datastructures import UploadFile
from litestar.enums import RequestEncodingType
from litestar.params import Body
from litestar.response import Template
from server.modules.attachments import AtachmentService
from server.modules.recognizer import RecognizerService
@dataclass
class BeerdsData:
f: UploadFile
device_id: str
class BreedsController(Controller):
path = "/beerds"
@post("/dogs")
async def beerds_dogs(
self,
request: Request,
data: Annotated[BeerdsData, Body(media_type=RequestEncodingType.MULTI_PART)],
) -> Response:
user_id = request.cookies.get("user_id")
# Cookie, которые нужно добавить в ответ (если пользователь новый)
new_cookies: dict[str, str] | None = None
if not user_id:
user_id = str(uuid.uuid4())
new_cookies = {"user_id": user_id}
recognizer_service: RecognizerService = inject.instance(RecognizerService)
body = await data.f.read()
result = await recognizer_service.predict_dog_image(body, user_id, data.device_id)
return Response(
content=result.to_serializable(),
media_type="application/json",
cookies=new_cookies,
)
@post("/cats")
async def beerds_cats(self, data: Annotated[UploadFile, Body(media_type=RequestEncodingType.MULTI_PART)]) -> dict:
recognizer_service: RecognizerService = inject.instance(RecognizerService)
body = await data.read()
result = await recognizer_service.predict_cat_image(body)
return result.to_serializable()
@get("/dogs/share/{result_id:str}")
async def beerds_share(self, result_id: str) -> Template:
recognizer_service: RecognizerService = inject.instance(RecognizerService)
result = await recognizer_service.get_results(result_id)
attach_service: AtachmentService = inject.instance(AtachmentService)
attachments = await attach_service.get_info_byid(session=None, attach_id=[result.attachment_id])
return Template(
template_name="share.html",
context={
"result": result,
"attachment": attachments[0],
},
)