Web 视觉应用:FastAPI + 图片风格转换

📂 所属阶段:第五阶段 — 工业落地与部署(实战篇)
🔗 相关章节:推理加速框架 · 边缘计算初探


1. 风格转换模型

import torch
import torchvision.transforms as transforms
from PIL import Image

class StyleTransfer:
    def __init__(self, model_path):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = torch.load(model_path).to(self.device)
        self.model.eval()
    
    def transfer(self, image_path):
        # 加载图像
        image = Image.open(image_path)
        
        # 预处理
        transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        x = transform(image).unsqueeze(0).to(self.device)
        
        # 推理
        with torch.no_grad():
            y = self.model(x)
        
        # 后处理
        y = y.squeeze(0).cpu()
        y = transforms.ToPILImage()(y)
        
        return y

2. FastAPI 服务

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
import io

app = FastAPI()
style_transfer = StyleTransfer("model.pth")

@app.post("/transfer")
async def transfer_style(file: UploadFile = File(...)):
    # 保存上传的文件
    contents = await file.read()
    image_path = f"/tmp/{file.filename}"
    with open(image_path, "wb") as f:
        f.write(contents)
    
    # 风格转换
    result = style_transfer.transfer(image_path)
    
    # 返回结果
    img_io = io.BytesIO()
    result.save(img_io, 'PNG')
    img_io.seek(0)
    
    return FileResponse(img_io, media_type="image/png")

# 运行:uvicorn app:app --host 0.0.0.0 --port 8000

3. 前端调用

<!DOCTYPE html>
<html>
<head>
    <title>Style Transfer</title>
</head>
<body>
    <input type="file" id="imageInput">
    <button onclick="transfer()">转换风格</button>
    <img id="result">
    
    <script>
        async function transfer() {
            const file = document.getElementById('imageInput').files[0];
            const formData = new FormData();
            formData.append('file', file);
            
            const response = await fetch('/transfer', {
                method: 'POST',
                body: formData
            });
            
            const blob = await response.blob();
            const url = URL.createObjectURL(blob);
            document.getElementById('result').src = url;
        }
    </script>
</body>
</html>

4. 小结

Web 视觉应用流程:

1. 模型:训练或加载预训练模型
2. API:用 FastAPI 包装模型
3. 前端:HTML + JavaScript 调用 API
4. 部署:Docker 容器化

完整流程:
- 用户上传图片
- 服务器处理
- 返回结果

💡 记住:FastAPI 是构建 AI 服务的最佳选择。它快速、易用、性能好。


🔗 扩展阅读