Update main.py

main.py

@@ -1,16 +1,74 @@
-import cv2,os
+import os
+import cv2
+from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Form, Request
+from fastapi.responses import HTMLResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+import shutil
 import AnimeGANv3_src
-if __name__ == '__main__':
 
-    f = 'A'
-    input_imgs_path = r'../../v3-usa\dataset\USA\val'
-    # input_imgs_path = r'/mnt/data/xinchen/v3-usa/dataset/USA/val'
-    output_path = 'AnimeGANv3_usa_64_output'
-    # img = cv2.imread(os.path.join(input_imgs_path, os.listdir(input_imgs_path)[0]))
-    img = cv2.imread(os.path.join(input_imgs_path, 'jp_16.png'))
-    out = AnimeGANv3_src.Convert(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), f, True)
-    # cv2.imshow('d', cv2.cvtColor(out, cv2.COLOR_BGR2RGB))
-    # cv2.waitKey(0)
-    cv2.imwrite('a.jpg', cv2.cvtColor(out, cv2.COLOR_BGR2RGB))
+app = FastAPI()
 
+os.makedirs('output', exist_ok=True)
 
+class InferenceRequest(BaseModel):
+    img_path: str
+    Style: str
+    if_face: str
+
+def inference(img_path, Style, if_face=None):
+    print(img_path, Style, if_face)
+    try:
+        img = cv2.imread(img_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        if Style == "AnimeGANv3_Arcane":
+            f = "A"
+        elif Style == "AnimeGANv3_Trump v1.0":
+            f = "T"
+        elif Style == "AnimeGANv3_Shinkai":
+            f = "S"
+        elif Style == "AnimeGANv3_PortraitSketch":
+            f = "P"
+        elif Style == "AnimeGANv3_Hayao":
+            f = "H"
+        elif Style == "AnimeGANv3_Disney v1.0":
+            f = "D"
+        elif Style == "AnimeGANv3_JP_face v1.0":
+            f = "J"
+        elif Style == "AnimeGANv3_Kpop v2.0":
+            f = "K"
+        else:
+            f = "U"
+
+        try:
+            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 RuntimeError as error:
+            print('Error', error)
+    except Exception as error:
+        print('global exception', error)
+        return None, None
+
+@app.post("/inference/")
+async def inference_api(request: InferenceRequest):
+    img_path = request.img_path
+    Style = request.Style
+    if_face = request.if_face
+    output, save_path = inference(img_path, Style, if_face)
+    return {"output": output, "save_path": save_path}
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")