65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, APIRouter
|
||
from fastapi.responses import Response
|
||
from starlette.responses import JSONResponse
|
||
from trustmark import TrustMark
|
||
from PIL import Image
|
||
import io
|
||
|
||
tm = TrustMark(model_type="Q", verbose=False)
|
||
appTmp = FastAPI(title="TrustMark Encode API", version="1.0")
|
||
app = APIRouter(prefix="/watermark")
|
||
|
||
@app.post("/encode")
|
||
async def encode_watermark(
|
||
image: UploadFile = File(...),
|
||
payload: str = Form(...),
|
||
outputType: str = Form(...),
|
||
strength: float = Form(1.2)
|
||
):
|
||
"""
|
||
给图片嵌入TrustMark盲水印
|
||
:param image: 原图文件上传
|
||
:param payload: 最多8个ASCII字符
|
||
:param outputType: 输出类型，必须是png，jpg，jpeg
|
||
:param strength: 取值一般0.0~1.5
|
||
:return: 带水印图片二进制
|
||
"""
|
||
img_bytes = await image.read()
|
||
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
||
watermarked_img = tm.encode(img, payload, WM_STRENGTH=strength)
|
||
buf = io.BytesIO()
|
||
watermarked_img.save(buf, format=outputType, quality=90)
|
||
buf.seek(0)
|
||
return Response(
|
||
content=buf.read(),
|
||
media_type=f"image/{outputType}"
|
||
)
|
||
|
||
@app.post("/decode", summary="上传图片提取TrustMark盲水印")
|
||
async def extract_watermark(image: UploadFile = File(...)):
|
||
MAX_FILE_SIZE = 10 * 1024 * 1024
|
||
if image.size and image.size > MAX_FILE_SIZE:
|
||
raise HTTPException(status_code=413, detail="文件过大，最大支持10MB")
|
||
allowed_types = ["image/png", "image/jpeg", "image/jpg"]
|
||
if image.content_type not in allowed_types:
|
||
raise HTTPException(status_code=400, detail="仅支持 png / jpg 图片")
|
||
|
||
try:
|
||
img_bytes = await image.read()
|
||
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
||
payload, has_watermark, schema = tm.decode(img)
|
||
|
||
return JSONResponse({
|
||
"success": True,
|
||
"has_watermark": has_watermark,
|
||
"payload": payload if has_watermark else None,
|
||
"schema": schema if has_watermark else None
|
||
})
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"解析失败: {str(e)}")
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
uvicorn.run("main:app", host="localhost", port=8081, workers=1)
|