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