Spaces:
Running
Running
from fastapi import FastAPI, File, UploadFile, Form | |
from fastapi.responses import StreamingResponse | |
from fastapi.middleware.cors import CORSMiddleware | |
from typing import List | |
import io | |
from PIL import Image, ImageOps | |
import numpy as np | |
import compColors | |
import dominantColors | |
import recolorReinhardAlgo | |
import recolorOTAlgo | |
import recolorTransferAlgo | |
import recolorLumaConverterAlgo | |
import recolorPaletteBasedTransfer | |
app = FastAPI() | |
app.add_middleware( | |
CORSMiddleware, | |
allow_origins=["*"], | |
allow_credentials=True, | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
async def dominant_color(file: UploadFile = File(...), num_colors: int = Form(...)): | |
""" | |
Receive an image file and an integer and return the dominant color(s). | |
""" | |
print('num_colors: ', num_colors) | |
file_content = await file.read() | |
image_bytes = io.BytesIO(file_content) | |
im = Image.open(image_bytes) | |
dominantColorsRGB = dominantColors.findDominantColors(image_bytes, num_colors, False) | |
dominantColorsHex = [dominantColors.rgb_to_hex(color) for color in dominantColorsRGB] | |
return {"dominantColors": dominantColorsHex} | |
async def color_palettes(colors: str = Form(...)): | |
""" | |
Receive an array of strings representing colors and return a color palette based on these colors. | |
""" | |
#maybe this isn't necessary. converting the string to an array of strings | |
colors = [color.strip() for color in colors.split(',')] | |
#generate the first pallete, which is the complementary colors of the given colors | |
complementaryColors = [] | |
for color in colors: | |
complementaryColors.append(compColors.complementary_colors(color)) | |
#generate the second palette using the adjacent colors algorithm: | |
adjacentColors = [] | |
for color in colors: | |
_adjcolors = compColors.adjacent_colors(color) | |
for _color in _adjcolors: | |
if _color not in adjacentColors: | |
adjacentColors.append(_color) | |
#generate the third palette using the gradient colors algorithm: | |
gradientColors = [] | |
for i in range(len(colors)-1): | |
gradientColors.append(compColors.gradient_colors(colors[i], colors[i+1])) | |
#Fixing size of palletes to 5 colors: | |
complementaryColors = [complementaryColors[i:i + 5] for i in range(0, len(complementaryColors), 5)] | |
adjacentColors = [adjacentColors[i:i + 5] for i in range(0, len(adjacentColors), 5)] | |
colors = [colors[i:i + 5] for i in range(0, len(colors), 5)] | |
return {"inputColor": colors, "complementaryColors": complementaryColors, "adjacentColors": adjacentColors, "gradientColors": gradientColors} | |
async def recolor(file: UploadFile = File(...), colors: str = Form(...), model: str = Form(...)): | |
""" | |
Receive an image file and an array of strings representing colors of a selected pallete and recolor an image. | |
""" | |
method = model | |
invertColors = False | |
colors = [color.strip() for color in colors.split(',')] | |
file_content = await file.read() | |
image_bytes = io.BytesIO(file_content) | |
image = Image.open(image_bytes) | |
if invertColors: | |
image = ImageOps.invert(image) | |
image_np = np.array(image) | |
if method == "CCA": | |
#Characteristic Color Analysis | |
recolorReinhardAlgo.recolor(image_np, colors) | |
elif method == "OTA": | |
#Optimal Transport Algorithm transfer | |
recolorOTAlgo.recolor(image_np, colors) | |
elif method =="KMEANS": | |
#K-means clustering transfer | |
recolorTransferAlgo.recolor(image_np, colors) | |
elif method == "LUMA": | |
#Luma converter transfer | |
recolorLumaConverterAlgo.remap_image_colors(image_np, colors) | |
elif method == "palette": | |
#palette transfer | |
recolorPaletteBasedTransfer.recolor(image_np, colors) | |
img_file = open("./result.jpg", "rb") | |
return StreamingResponse(img_file, media_type="image/jpeg") | |
async def test(): | |
""" | |
Test endpoint to check if the server is running. | |
""" | |
return {"message": "Server is running"} | |
if __name__ == "__main__": | |
import uvicorn | |
print("Server is running") | |
uvicorn.run(app, host="0.0.0.0", port=7860) | |
#how to run: | |
#source env/bin/activate | |
#uvicorn server:app --reload |