Spaces:

99i
/

h2c

Sleeping

h2c

File size: 14,653 Bytes

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 """
    <!DOCTYPE html>
    <html>
    <head>
        <title>HTML to Image Converter</title>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <style>
            body {
                font-family: Arial, sans-serif;
                max-width: 800px;
                margin: 0 auto;
                padding: 20px;
                background-color: #f5f5f5;
            }
            .container {
                background-color: white;
                padding: 30px;
                border-radius: 10px;
                box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            }
            h1 {
                color: #333;
                text-align: center;
            }
            .form-group {
                margin-bottom: 20px;
            }
            label {
                display: block;
                margin-bottom: 5px;
                font-weight: bold;
            }
            input, textarea, select {
                width: 100%;
                padding: 10px;
                border: 1px solid #ddd;
                border-radius: 5px;
                box-sizing: border-box;
            }
            button {
                background-color: #4CAF50;
                color: white;
                padding: 12px 20px;
                border: none;
                border-radius: 5px;
                cursor: pointer;
                font-size: 16px;
                width: 100%;
            }
            button:hover {
                background-color: #45a049;
            }
            #result {
                margin-top: 20px;
                text-align: center;
            }
            #imagePreview {
                max-width: 100%;
                border: 1px solid #ddd;
                border-radius: 5px;
                margin-top: 10px;
            }
            .tab {
                overflow: hidden;
                border: 1px solid #ccc;
                background-color: #f1f1f1;
            }
            .tab button {
                background-color: inherit;
                float: left;
                border: none;
                outline: none;
                cursor: pointer;
                padding: 14px 16px;
                transition: 0.3s;
                font-size: 17px;
            }
            .tab button:hover {
                background-color: #ddd;
            }
            .tab button.active {
                background-color: #ccc;
            }
            .tabcontent {
                display: none;
                padding: 6px 12px;
                border: 1px solid #ccc;
                border-top: none;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>HTML to Image Converter</h1>
            
            <div class="tab">
              <button class="tablinks active" onclick="openTab(event, 'urlTab')">URL地址</button>
              <button class="tablinks" onclick="openTab(event, 'htmlTab')">HTML代码</button>
            </div>

            <div id="urlTab" class="tabcontent" style="display: block;">
                <form id="imageFormUrl">
                    <div class="form-group">
                        <label for="url">网页URL:</label>
                        <input type="url" id="url" name="url" required placeholder="https://example.com">
                    </div>
                    <div class="form-group">
                        <label for="selector">CSS选择器 (可选):</label>
                        <input type="text" id="selector" name="selector" placeholder="例如: .content, #main">
                    </div>
                    <div class="form-group">
                        <label for="full_page">截图类型:</label>
                        <select id="full_page" name="full_page">
                            <option value="false">可见区域</option>
                            <option value="true">完整页面</option>
                        </select>
                    </div>
                    <div class="form-group">
                        <label for="width">宽度:</label>
                        <input type="number" id="width" name="width" value="1200" min="100" max="5000">
                    </div>
                    <div class="form-group">
                        <label for="height">高度:</label>
                        <input type="number" id="height" name="height" value="800" min="100" max="5000">
                    </div>
                    <button type="submit">生成图片</button>
                </form>
            </div>

            <div id="htmlTab" class="tabcontent">
                <form id="imageFormHtml">
                    <div class="form-group">
                        <label for="html_content">HTML代码:</label>
                        <textarea id="html_content" name="html_content" rows="10" placeholder="在这里输入HTML代码..."></textarea>
                    </div>
                    <div class="form-group">
                        <label for="selectorHtml">CSS选择器 (可选):</label>
                        <input type="text" id="selectorHtml" name="selector" placeholder="例如: .content, #main">
                    </div>
                    <div class="form-group">
                        <label for="full_page_html">截图类型:</label>
                        <select id="full_page_html" name="full_page">
                            <option value="false">可见区域</option>
                            <option value="true">完整页面</option>
                        </select>
                    </div>
                    <div class="form-group">
                        <label for="width_html">宽度:</label>
                        <input type="number" id="width_html" name="width" value="1200" min="100" max="5000">
                    </div>
                    <div class="form-group">
                        <label for="height_html">高度:</label>
                        <input type="number" id="height_html" name="height" value="800" min="100" max="5000">
                    </div>
                    <button type="submit">生成图片</button>
                </form>
            </div>
            
            <div id="result"></div>
        </div>
        
        <script>
            // 切换标签页
            function openTab(evt, tabName) {
                var i, tabcontent, tablinks;
                tabcontent = document.querySelectorAll(".tabcontent");
                for (i = 0; i < tabcontent.length; i++) {
                    tabcontent[i].style.display = "none";
                }
                tablinks = document.querySelectorAll(".tablinks");
                for (i = 0; i < tablinks.length; i++) {
                    tablinks[i].className = tablinks[i].className.replace(" active", "");
                }
                document.getElementById(tabName).style.display = "block";
                evt.currentTarget.className += " active";
            }
            
            // URL表单提交处理
            document.getElementById('imageFormUrl').addEventListener('submit', async function(e) {
                e.preventDefault();
                const resultDiv = document.getElementById('result');
                resultDiv.innerHTML = '<p>正在处理...</p>';
                
                const formData = new FormData(this);
                const data = Object.fromEntries(formData);
                data['html_content'] = null; // 清除HTML内容
                
                try {
                    const response = await fetch('/convert', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify(data)
                    });
                    
                    const result = await response.json();
                    
                    if (result.success) {
                        resultDiv.innerHTML = `
                            <h2>转换成功!</h2>
                            <img id="imagePreview" src="${result.image_url}" alt="生成的图片">
                        `;
                    } else {
                        resultDiv.innerHTML = `<p style="color: red;">错误: ${result.error}</p>`;
                    }
                } catch (error) {
                    resultDiv.innerHTML = `<p style="color: red;">请求失败: ${error.message}</p>`;
                }
            });
            
            // HTML表单提交处理
            document.getElementById('imageFormHtml').addEventListener('submit', async function(e) {
                e.preventDefault();
                const resultDiv = document.getElementById('result');
                resultDiv.innerHTML = '<p>正在处理...</p>';
                
                const formData = new FormData(this);
                const data = Object.fromEntries(formData);
                data['url'] = null; // 清除URL
                
                try {
                    const response = await fetch('/convert', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json',
                        },
                        body: JSON.stringify(data)
                    });
                    
                    const result = await response.json();
                    
                    if (result.success) {
                        resultDiv.innerHTML = `
                            <h2>转换成功!</h2>
                            <img id="imagePreview" src="${result.image_url}" alt="生成的图片">
                        `;
                    } else {
                        resultDiv.innerHTML = `<p style="color: red;">错误: ${result.error}</p>`;
                    }
                } catch (error) {
                    resultDiv.innerHTML = `<p style="color: red;">请求失败: ${error.message}</p>`;
                }
            });
        </script>
    </body>
    </html>
    """

@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"""
<html>
                            <head></head>
                            <body id="ui" style="width:{request.width}px;height:{request.height}px;">{request.html_content}</body>
</html>
""")
                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)