Spaces:
Sleeping
Sleeping
File size: 2,818 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 |
import os
from pyppeteer import launch
from PIL import Image
import json
import random
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
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)}"
}) |