main.py

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))