原则:
- 输入输出都基于BaseModel
- 依靠JSONResponse制定返回错误的json信息
- 依靠装饰器中@app.post制定responses字典从而让docs文档更丰富
import uvicorn
from pydantic import BaseModel, Field
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import BlipProcessor, BlipForConditionalGeneration
from fastapi.responses import JSONResponse
from PIL import Image
import io
import base64
class UpscalerRequest(BaseModel):
base64_image: str = Field(
...,
title="Base64 Encoded Image",
description="The base64-encoded image that you want to upscale."
)
outscale: float = Field(
...,
ge=1.0, # 大于等于1.0
le=5.0, # 小于等于5.0
title="Upscale Factor",
description="The scaling factor for image upscaling. Should be between 1.0 and 5.0."
)
class UpscalerResponse(BaseModel):
base64_image_out: str = Field(
...,
title="Base64 Encoded Upscaled Image",
description="The base64-encoded image after upscaling."
)
class CustomErrorResponse:
def __init__(self, description: str, error_code: int, detail: str):
self.description = description
self.error_code = error_code
self.detail = detail
def to_response_dict(self):
return {
"description": self.description,
"content": {
"application/json": {
"example": {"error_code": self.error_code, "detail": self.detail}
}
},
}
def __call__(self, extra_detail=None):
if extra_detail is not None:
self.detail = f"detail:{self.detail}, extra_detail:{extra_detail}"
return JSONResponse(content={"error_code": self.error_code, "detail": self.detail}, status_code=self.error_code)
image_upscaler_CustomErrorResponse = CustomErrorResponse("超分执行错误", 501, "upscale error")
@app.post("/image_upscaler", response_model=UpscalerResponse, summary="Image Upscaler",
responses={image_upscaler_CustomErrorResponse.error_code: image_upscaler_CustomErrorResponse.to_response_dict()})
def image_upscaler(request: UpscalerRequest):
"""
Image Upscaler.
Parameters:
- `base64_image`: The base64-encoded image.
- `outscale`: The scaling factor for image upscaling (between 1.0 and 5.0).
Returns:
- `output`: The base64-encoded upscaled image.
Example:
```python
import requests
import base64
from PIL import Image
from io import BytesIO
# 1. 读取图像文件并转换为base64字符串
image_path = "car.png"
with open(image_path, "rb") as image_file:
base64_image = base64.b64encode(image_file.read()).decode("utf-8")
# 2. 构造请求数据
outpainting_request = {
"base64_image": base64_image,
"outscale": 3,
}
# 3. 发送HTTP POST请求
api_url = "http://home.elvisiky.com:7862/image_upscaler"
response = requests.post(api_url, json=outpainting_request)
# 4. 处理响应
if response.status_code == 200:
result_data = response.json()
# 5. 保存base64编码的图像为文件
detected_map_base64 = result_data["base64_image_out"]
detected_map_image = Image.open(BytesIO(base64.b64decode(detected_map_base64)))
detected_map_image.save("base64_image_out.png")
print("图像保存成功!")
else:
print(f"API调用失败,HTTP状态码:{response.status_code}")
print(response.json())
```
"""
try:
base64_image = request.base64_image
logging.info(f"image_upscaler, image length {len(base64_image)}")
img = decode_image_from_base64(base64_image)
img_cv2 = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
output = enhance_image(img_cv2, outscale=request.outscale)
if output is None:
return image_upscaler_CustomErrorResponse()
else:
output_base64 = base64.b64encode(cv2.imencode('.png', output)[1]).decode('utf-8')
return {"base64_image_out": output_base64}
except Exception as e:
return image_upscaler_CustomErrorResponse(str(e))
if __name__ == '__main__':
uvicorn.run(f'{os.path.basename(__file__).split(".")[0]}:app',
host='0.0.0.0',
port=7862,
reload=False,
workers=1)