rest_api.py - master - IvlevAlexander0/QNN_FastApi

rest_api.py
113 строк4 KB
IvlevAlexander0
ver0.9.1
23 окт 2025, 08:17
23 окт 2025, 08:170d124a2
import numpy as np
from PIL import Image
import io
import base64
import warnings
warnings.filterwarnings('ignore')

from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import HTMLResponse, Response
from fastapi.templating import Jinja2Templates

from torch.cuda import is_available
from classification.classifier import return_classifier
from segmentation.segmenter import return_segmenter


classifier_config = {
    'device': 'cuda' if is_available() else 'cpu',   # Устройство исполнения (cuda или cpu)
    'image_size': 224,  # Размер изoбражения подаваемого на вход модели (3ximage_sizeximage_size);
    'classifier_model': "./classification/classifier_model.pth",# Путь к сохранённой модели
}

classifier = return_classifier(classifier_config)

segmenter_config = {
    'device': 'cuda' if is_available() else 'cpu',   # Устройство исполнения (cuda или cpu)
    'target_size': (224, 224), # Размер изображения которое пойдет на вход в модель
    'checkpoint': "./segmentation/segmenter_model.pth", # Путь к сохраненной модели
    'transform_type': "resize" # 'resize' - подгонка размера с сохранением соотношения сторон
}

segmenter = return_segmenter(segmenter_config)

app = FastAPI(title="Image Processing API")
changed_url = "/{process}/ports/{ports}/"

templates = Jinja2Templates(directory="templates")

def image_to_base64(image: Image.Image, format: str) -> str:
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format=format)
    img_byte_arr = img_byte_arr.getvalue()
    return base64.b64encode(img_byte_arr).decode('utf-8')

def get_base_url(request: Request):
    return str(request.url).rstrip('/')

@app.get(changed_url, response_class=HTMLResponse)
async def read_root(request: Request):
    base_url = get_base_url(request)
    return templates.TemplateResponse("upload.html", {"request": request, "base_url": base_url})

@app.post(changed_url + "upload/", response_class=HTMLResponse)
async def upload_file(request: Request, file: UploadFile = File(...)):
    try:
        base_url = get_base_url(request)[:-7]

        if not file.content_type.startswith('image/'):
            return templates.TemplateResponse(
                "upload.html", 
                {"request": request, "error": "Файл должен быть изображением", "base_url": base_url}
            )

        image_data = await file.read()
        image = Image.open(io.BytesIO(image_data))

        original_format = image.format or "JPEG"
        original_base64 = image_to_base64(image, original_format)

        image_class = classifier(image)
        segmentation_image = segmenter(image)
        segmentation_image = (segmentation_image * 255).astype(np.uint8)

        segmentation_base64 = image_to_base64(Image.fromarray(segmentation_image).convert('RGB'), format=original_format)

        mime_type = f"image/{original_format.lower()}"

        return templates.TemplateResponse(
            "result.html", 
            {
                "request": request, 
                "original_image": f"data:{mime_type};base64,{original_base64}",
                "segmentation_image": f"data:{mime_type};base64,{segmentation_base64}",
                "message": "Изображение успешно обработано",
                "class": "Класс изображения: " + image_class,
                "base_url": base_url,
                "filename": file.filename
            }
        )

    except Exception as e:
        return templates.TemplateResponse(
            "upload.html", 
            {"request": request, "error": f"Ошибка обработки: {str(e)}", "base_url": base_url}
        )

@app.post(changed_url + "download/")
async def download_file(request: Request):
    form_data = await request.form()
    image_data = form_data["image_data"]
    filename = form_data["filename"]

    header, base64_str = image_data.split(",", 1)

    image_bytes = base64.b64decode(base64_str)

    content_type = header.split(";")[0].split(":")[1]

    return Response(
        content=image_bytes,
        media_type=content_type,
        headers={"Content-Disposition": f"attachment; filename=segmentation_{filename}"}
    )