#Web Vision Application: Detailed explanation of FastAPI, image processing and AI service deployment

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")