Update main.py

main.py

@@ -1,16 +1,64 @@
-import cv2,os
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+from fastapi.responses import StreamingResponse
+import os
+import cv2
+import numpy as np
 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))
+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))