from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional
import asyncio
import tempfile
import os
import uuid
from DrissionPage import ChromiumPage,ChromiumOptions
co = ChromiumOptions()
co.set_argument('--no-sandbox').headless()
app = FastAPI(title="HTML to Image Converter", version="1.0.0")
# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")
# 创建临时目录用于存储图片
os.makedirs("temp", exist_ok=True)
class ImageRequest(BaseModel):
url: Optional[str] = None
html_content: Optional[str] = None
selector: Optional[str] = None
full_page: bool = False
width: int = 1200
height: int = 800
class ImageResponse(BaseModel):
success: bool
image_url: Optional[str] = None
error: Optional[str] = None
@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
return """
HTML to Image Converter
HTML to Image Converter
"""
@app.post("/convert", response_model=ImageResponse)
async def convert_html_to_image(request: ImageRequest):
try:
# 生成唯一的临时文件名
unique_id = str(uuid.uuid4())
temp_image_path = f"temp/{unique_id}.png"
# 使用DrissionPage进行截图
from DrissionPage import ChromiumPage
import os
import time
# 简单的重试逻辑避免浏览器启动问题
max_retries = 3
retry_count = 0
page = None
while retry_count < max_retries:
try:
# 创建页面实例（使用简单的参数避免环境兼容性问题）
page = ChromiumPage(co)
break
except Exception as e:
retry_count += 1
if retry_count >= max_retries:
raise e
time.sleep(0.5) # 等待后重试
if page is None:
raise HTTPException(status_code=500, detail="浏览器启动失败，请检查环境配置")
try:
page.set
# 根据是URL还是HTML内容来处理
if request.url:
# 打开URL
page.get(request.url)
# 等待页面加载
page.wait.loadstart()
page.wait.loadfinish()
elif request.html_content:
# 使用HTML内容
with open("tem.html","w+",encoding="utf-8") as o:
o.write(f"""
{request.html_content}
""")
page.get("tem.html")
else:
raise HTTPException(status_code=400, detail="必须提供URL或HTML内容")
# 根据参数决定截图方式
if request.selector:
element = page.ele(request.selector)
if element:
element.screenshot(path=temp_image_path)
else:
raise HTTPException(status_code=400, detail="CSS选择器未找到元素")
elif request.full_page:
page.screenshot(path=temp_image_path, full_page=True)
else:
page("#ui").get_screenshot(path="static/new.png")
finally:
if page:
page.quit()
return ImageResponse(
success=True,
image_url=f"static/new.png"
)
except Exception as e:
return ImageResponse(
success=False,
error=str(e)
)
@app.get("/{image_path:path}")
async def get_image(image_path: str):
# 安全检查，只允许访问temp目录下的文件
if not image_path.startswith("temp/"):
raise HTTPException(status_code=404, detail="Image not found")
# 检查文件是否存在
if not os.path.exists(image_path):
raise HTTPException(status_code=404, detail="Image not found")
# 返回图片文件
return FileResponse(image_path)
@app.get("/health")
async def health_check():
return {"status": "healthy"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)