#Web视觉应用：FastAPI、图像处理与AI服务部署详解

from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
from typing import Optional
from datetime import datetime
import uuid

# 初始化 FastAPI 应用
app = FastAPI(
    title="AI 风格迁移 API",
    description="极简风格迁移Web服务",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# 统一的响应格式（后续所有接口复用）
class StandardResponse(BaseModel):
    success: bool
    message: str
    data: Optional[dict] = None
    timestamp: datetime = datetime.now()
    request_id: str = str(uuid.uuid4())

# 根路径
@app.get("/", response_model=StandardResponse)
async def root():
    return StandardResponse(
        success=True,
        message="欢迎使用AI视觉服务，请访问 /docs 查看API文档",
        data={"api_version": "1.0.0"}
    )

# 临时图像上传验证接口
@app.post("/api/v1/upload", response_model=StandardResponse)
async def upload_temp_image(file: UploadFile = File(...)):
    # 1. 检查文件大小（限制 10MB）
    file_size = len(await file.read())
    if file_size > 10 * 1024 * 1024:
        raise HTTPException(status_code=400, detail="文件大小不能超过10MB")
    # 2. 重置文件指针，以便后续读取
    await file.seek(0)
    # 3. 检查文件类型
    allowed_mime = ["image/jpeg", "image/png", "image/jpg"]
    if file.content_type not in allowed_mime:
        raise HTTPException(status_code=400, detail="仅支持JPEG/PNG格式图像")
    
    return StandardResponse(
        success=True,
        message="临时图像上传成功",
        data={"filename": file.filename, "size": file_size}
    )

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import os
import threading
from typing import Optional

class VisionModel:
    """视觉模型服务基类"""
    def __init__(
        self,
        model_path: str,
        device: Optional[str] = None,
        input_size: tuple = (224, 224),
        mean: list = [0.485, 0.456, 0.406],
        std: list = [0.229, 0.224, 0.225]
    ):
        # 1. 自动选择设备（优先 GPU）
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        # 2. 模型加载锁（防止多线程同时加载/推理冲突）
        self._model_lock = threading.Lock()
        # 3. 预处理与后处理参数
        self.input_size = input_size
        self.mean = mean
        self.std = std
        # 4. 加载模型并设为推理模式
        self.model = self._load_model(model_path)
        self.model.eval()
        # 5. 构建预处理与后处理变换
        self.preprocess_transform = self._build_preprocess()
        self.postprocess_transform = self._build_postprocess()

    def _load_model(self, model_path: str):
        """加载 PyTorch 模型（本地测试用，生产环境建议转 ONNX/TensorRT）"""
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        with self._model_lock:
            model = torch.load(model_path, map_location=self.device)
        return model.to(self.device)

    def _build_preprocess(self):
        """构建图像预处理变换"""
        return transforms.Compose([
            transforms.Resize(self.input_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=self.mean, std=self.std)
        ])

    def _build_postprocess(self):
        """构建图像后处理变换（风格迁移用，分类/检测需要重写）"""
        mean_inv = [-m/s for m, s in zip(self.mean, self.std)]
        std_inv = [1/s for s in self.std]
        return transforms.Compose([
            transforms.Normalize(mean=mean_inv, std=std_inv),
            transforms.Lambda(lambda x: torch.clamp(x, 0, 1)),
            transforms.ToPILImage()
        ])

    def preprocess(self, image: Image.Image) -> torch.Tensor:
        """预处理单张图像"""
        return self.preprocess_transform(image).unsqueeze(0).to(self.device)

    def postprocess(self, output: torch.Tensor) -> Image.Image:
        """后处理单张风格迁移结果"""
        output = output.squeeze(0).cpu()
        return self.postprocess_transform(output)

    def infer(self, image: Image.Image) -> Image.Image:
        """加锁推理（多线程安全）"""
        with torch.no_grad():          # 关闭梯度计算，节省显存
            with self._model_lock:
                input_tensor = self.preprocess(image)
                output_tensor = self.model(input_tensor)
        return self.postprocess(output_tensor)

# 实例化风格迁移模型（使用更大的输入尺寸）
try:
    style_model = VisionModel(
        model_path="models/starry_night.pth",
        input_size=(512, 512)  # 风格迁移建议用更大分辨率
    )
    print(f"✅ 风格迁移模型加载成功，运行设备: {style_model.device}")
except Exception as e:
    print(f"❌ 模型加载失败: {str(e)}")
    style_model = None

from fastapi.responses import FileResponse
import shutil
import tempfile

# ...（前面的 FastAPI 初始化、StandardResponse、VisionModel 代码）

@app.post("/api/v1/style-transfer", response_model=StandardResponse)
async def style_transfer(file: UploadFile = File(...)):
    if not style_model:
        raise HTTPException(status_code=500, detail="模型未加载成功，请联系管理员")

    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    temp_input = os.path.join(temp_dir, file.filename)
    temp_output = os.path.join(temp_dir, f"result_{uuid.uuid4().hex[:8]}.jpg")

    try:
        # 1. 保存用户上传的原始图片
        with open(temp_input, "wb") as f:
            shutil.copyfileobj(file.file, f)
        
        # 2. 读取图片并进行推理
        image = Image.open(temp_input).convert("RGB")
        result_img = style_model.infer(image)
        
        # 3. 保存结果
        result_img.save(temp_output, quality=95)
        
        # 4. 返回结果文件（响应完成后自动清理临时目录）
        return FileResponse(
            path=temp_output,
            filename=f"starry_night_{file.filename}",
            media_type="image/jpeg",
            background=shutil.rmtree(temp_dir)  # 后台清理
        )

    except Exception as e:
        shutil.rmtree(temp_dir)   # 出错也要清理
        raise HTTPException(status_code=500, detail=f"推理失败: {str(e)}")

from fastapi.staticfiles import StaticFiles

# 挂载静态文件目录（访问 http://localhost:8000/ 即可看到前端页面）
app.mount("/", StaticFiles(directory="static", html=True), name="static")