farm-recolor / server.py
vettorazi's picture
fixing docker to accept saving files
16dea1d
raw
history blame
No virus
4.31 kB
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=["*"],
)
@app.post("/dominantColor/")
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}
@app.post("/ColorPalettes/")
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}
@app.post("/recolor/")
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")
@app.get("/test/")
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