File size: 2,158 Bytes
0860266 b117429 0860266 b117429 0860266 b117429 |
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 |
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import StreamingResponse
import os
import cv2
import numpy as np
import AnimeGANv3_src
from io import BytesIO
# Initialize FastAPI
app = FastAPI()
os.makedirs('output', exist_ok=True)
def process_image(img_path, style, if_face):
print(img_path, style, if_face)
try:
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
style_mapping = {
"AnimeGANv3_Arcane": "A",
"AnimeGANv3_Trump v1.0": "T",
"AnimeGANv3_Shinkai": "S",
"AnimeGANv3_PortraitSketch": "P",
"AnimeGANv3_Hayao": "H",
"AnimeGANv3_Disney v1.0": "D",
"AnimeGANv3_JP_face v1.0": "J",
"AnimeGANv3_Kpop v2.0": "K",
"AnimeGANv3_USA": "U"
}
f = style_mapping.get(style, "U")
det_face = True if if_face == "Yes" else False
output = AnimeGANv3_src.Convert(img, f, det_face)
save_path = f"output/out.{img_path.rsplit('.')[-1]}"
cv2.imwrite(save_path, output[:, :, ::-1])
return output, save_path
except Exception as error:
print('Error', error)
return None, None
@app.post("/inference/")
async def inference(file: UploadFile = File(...), Style: str = Form(...), if_face: str = Form(...)):
try:
# Save the uploaded file to a temporary location
file_location = f"temp_{file.filename}"
with open(file_location, "wb") as f:
f.write(await file.read())
# Process the image
output, save_path = process_image(file_location, Style, if_face)
if output is None:
raise HTTPException(status_code=500, detail="Processing failed")
# Read the processed image and prepare it for response
with open(save_path, "rb") as result_file:
result_bytes = result_file.read()
# Return the image as a blob
return StreamingResponse(BytesIO(result_bytes), media_type="image/jpeg")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|