Spaces:
Sleeping
Sleeping
File size: 2,867 Bytes
aef14d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import os
from pyppeteer import launch
from PIL import Image
import json
import random
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
from jinja2 import Environment, FileSystemLoader
import io
from fastapi.responses import StreamingResponse
app = FastAPI()
class HtmlInput(BaseModel):
message: str
async def convert_html_to_image(html, output_file, width=650, height=800, device_scale_factor=2):
try:
browser = await launch({
'headless': True,
'args': ['--no-sandbox', '--disable-setuid-sandbox'],
'timeout': 60000 # Increase timeout if needed
})
page = await browser.newPage()
await page.setViewport({'width': width, 'height': height, 'deviceScaleFactor': device_scale_factor})
await page.setContent(html)
await page.screenshot({'path': output_file})
await browser.close()
return output_file # Return the path of the captured image
except Exception as e:
print(f"Error in convert_html_to_image: {e}")
raise
async def optimize_image(input_file, output_file, target_size_kb=700):
# Open the image using PIL
with Image.open(input_file) as img:
# Calculate the initial quality to achieve the target size
quality = 95
temp_output_file = output_file + ".temp.jpg"
img.save(temp_output_file, format='JPEG', quality=quality, optimize=True)
while os.path.getsize(temp_output_file) > target_size_kb * 1024 and quality > 0:
# Reduce the quality and save to temporary file
quality -= 5
img.save(temp_output_file, format='JPEG', quality=quality, optimize=True)
# Save the optimized image to the output file
os.replace(temp_output_file, output_file)
async def mainFunction(html_content:str):
output_image_file = "image.jpg"
captured_image_path = await convert_html_to_image(html_content, output_image_file)
# Optimize the image to be less than 200KB
await optimize_image(captured_image_path, output_image_file)
return output_image_file
# Run the asynchronous main function
@app.post("/convert-html-to-image/")
async def convert_html_to_image_endpoint(html_content:HtmlInput):
image_path = await mainFunction(html_content.message)
# Check if image was generated
if not image_path or not os.path.exists(image_path):
raise HTTPException(status_code=500, detail="Image generation failed.")
# Open the image file for streaming
with open(image_path, "rb") as img_file:
buffer = io.BytesIO(img_file.read())
buffer.seek(0)
# Return the image as a downloadable file
return StreamingResponse(buffer, media_type="image/jpeg", headers={
"Content-Disposition": f"attachment; filename={os.path.basename(image_path)}"
}) |