144 lines
4.5 KiB
Python
144 lines
4.5 KiB
Python
import datetime
|
|
|
|
from fastapi import FastAPI, File, UploadFile
|
|
from PIL import Image
|
|
import os
|
|
import io
|
|
from typing import List
|
|
from pydantic import BaseModel
|
|
import random
|
|
import glob
|
|
import hashlib
|
|
|
|
app = FastAPI()
|
|
|
|
BACKGROUND_DIR = 'backgrounds'
|
|
IMAGE_DIR = 'images'
|
|
IMAGE_EXT = '.png'
|
|
|
|
|
|
class MosaicRequest(BaseModel):
|
|
rows: int
|
|
columns: int
|
|
buffer_size: int
|
|
keywords: List[str]
|
|
background_file: str
|
|
|
|
class BackgroundRequest(BaseModel):
|
|
backgrounds: List[str]
|
|
|
|
|
|
def get_images():
|
|
image_files = [f for f in os.listdir(IMAGE_DIR) if f.endswith(IMAGE_EXT)]
|
|
image_files.sort()
|
|
return [Image.open(os.path.join(IMAGE_DIR, f)) for f in image_files]
|
|
|
|
|
|
def create_mosaic(images, rows, cols, buffer_size):
|
|
mosaic_width = cols * images[0].size[0] + (cols - 1) * buffer_size
|
|
mosaic_height = rows * images[0].size[1] + (rows - 1) * buffer_size
|
|
mosaic = Image.new('RGBA', (mosaic_width, mosaic_height), color=(0, 0, 0, 0))
|
|
for i, img in enumerate(images):
|
|
row = i // cols
|
|
col = i % cols
|
|
offset_x = (images[0].size[0] - img.size[0]) // 2
|
|
offset_y = (images[0].size[1] - img.size[1]) // 2
|
|
pos_x = col * (img.size[0] + buffer_size)
|
|
pos_y = row * (img.size[1] + buffer_size)
|
|
pos = (pos_x + offset_x, pos_y + offset_y)
|
|
mosaic.paste(img, pos)
|
|
return mosaic
|
|
|
|
|
|
def place_image_on_background(foreground: Image.Image, background: Image.Image) -> Image.Image:
|
|
# Compute the position to paste the foreground image
|
|
x = (background.width - foreground.width) // 2
|
|
y = (background.height - foreground.height) // 2
|
|
|
|
# Create a new image with the background and paste the foreground onto it
|
|
result = Image.new("RGBA", (background.width, background.height), (255, 255, 255, 0))
|
|
result.paste(background, (0, 0))
|
|
result.alpha_composite(foreground, (x, y))
|
|
|
|
return result
|
|
|
|
@app.get("/backgrounds/")
|
|
async def list_backgrounds():
|
|
background_files = [f for f in os.listdir(BACKGROUND_DIR) if f.endswith(IMAGE_EXT)]
|
|
background_files.sort()
|
|
background_dict = {"backgrounds": background_files}
|
|
return BackgroundRequest(**background_dict)
|
|
|
|
@app.post("/upload_file/")
|
|
async def upload_file(file: UploadFile = File(...), directory=IMAGE_DIR):
|
|
# check if the file is an image file
|
|
if not file.content_type.startswith('image'):
|
|
raise ValueError('File is not an image')
|
|
|
|
# read the file data and create a Pillow Image object
|
|
image = Image.open(file.file)
|
|
|
|
# save the PNG image to the "images/" directory
|
|
filename = os.path.join(directory, os.path.splitext(file.filename)[0] + '.png')
|
|
image.save(filename)
|
|
|
|
return {"filename": filename}
|
|
|
|
|
|
@app.post("/upload_background")
|
|
async def upload_background(file: UploadFile = File(...)):
|
|
filename_d = await upload_file(file, directory=BACKGROUND_DIR)
|
|
return filename_d
|
|
|
|
|
|
@app.post("/generate_mosaic")
|
|
async def generate_mosaic(request: MosaicRequest):
|
|
images = []
|
|
for keyword in request.keywords:
|
|
file_list = glob.glob(f'images/{keyword}*.png')
|
|
if not file_list:
|
|
return {"message": f"No images found for keyword: {keyword}"}
|
|
selected_file = random.choice(file_list)
|
|
with open(selected_file, "rb") as f:
|
|
images.append(Image.open(io.BytesIO(f.read())))
|
|
|
|
# select background image
|
|
background = Image.open(f"backgrounds/{request.background_file}")
|
|
|
|
# resize images to have the same height
|
|
height = min(background.height - (request.rows + 1) * request.buffer_size,
|
|
min(img.size[1] for img in images)) / request.rows
|
|
images = [img.resize((round(img.size[0]*height/img.size[1]),
|
|
round(height))) for img in images]
|
|
|
|
# create grid
|
|
grid = create_mosaic(images, request.rows, request.columns, request.buffer_size)
|
|
|
|
# select background image
|
|
background = Image.open(f"backgrounds/{request.background_file}")
|
|
|
|
# place grid on background
|
|
final_image = place_image_on_background(grid, background)
|
|
|
|
# convert image to bytes
|
|
img_bytes = final_image.tobytes()
|
|
|
|
# create MD5 hash object
|
|
hash_obj = hashlib.md5()
|
|
|
|
# update hash object with image bytes
|
|
hash_obj.update(img_bytes)
|
|
|
|
# get hex representation of hash
|
|
hash_hex = hash_obj.hexdigest()
|
|
|
|
# convert hex to string
|
|
hash_str = str(hash_hex)
|
|
|
|
# save final image
|
|
now = datetime.datetime.now()
|
|
filename_datetime = now.strftime("%Y-%m-%d_%H-%M-%S")
|
|
|
|
final_image.save(f"mosaics/{hash_str}-{filename_datetime}.png")
|
|
|
|
return {"message": "Mosaic image created successfully"}
|