Update main.py

main.py

@@ -5,12 +5,12 @@ from base64 import b64encode
 from speech_recognition import AudioFile, Recognizer
 import numpy as np
 from scipy.spatial import distance as dist
-from typing import Union
+from typing import Union, Optional
 
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse
 
-from utils import tts, read_image_file, pil_to_base64, get_hist
+from utils import tts, read_image_file, pil_to_base64, base64_to_pil, get_hist
 from typing import Optional
 from huggingface_hub import hf_hub_download
 
@@ -24,7 +24,7 @@ model = YOLO(model_path)
 CLASS = model.model.names
 defaul_bot_voice = "おはいようございます"
 area_thres = 0.3
-
+ZIP = False
 app = FastAPI()
 
 @app.get("/")
@@ -35,7 +35,8 @@ def read_root():
 @app.post("/aisatsu_api/")
 async def predict_api(
         file: UploadFile = File(...),
-        last_seen: Union[UploadFile, None] = File(None)
+        # last_seen: Union[UploadFile, None] = File(None),
+        last_seen: Optional[str] = Form(None),
 ):
     total_time = time.time()
     start_time = time.time()
@@ -62,7 +63,10 @@ async def predict_api(
     print("Get face time", time.time() - start_time)
     start_time = time.time()
     if last_seen is not None:
-        last_seen = read_image_file(await last_seen.read())
+        if type(last_seen) == str:
+            last_seen = base64_to_pil(last_seen)
+        else:
+            last_seen = read_image_file(await last_seen.read())
         if out_img is not None:
             diff_value = dist.euclidean(get_hist(out_img), get_hist(last_seen))
     print("Hist time", time.time() - start_time)
@@ -72,17 +76,24 @@ async def predict_api(
         voice_bot_path = tts(defaul_bot_voice, language="ja")
         image_bot_path = pil_to_base64(out_img)
         print("Voice time", time.time() - start_time)
-        io = BytesIO()
-        zip_filename = "final_archive.zip"
-        with zipfile.ZipFile(io, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
-            for file_path in [voice_bot_path, image_bot_path]:
-                zf.write(file_path)
-            zf.close()
-        print("Total time", time.time() - total_time)
-        return StreamingResponse(
-            iter([io.getvalue()]),
-            media_type="application/x-zip-compressed",
-            headers={"Content-Disposition": f"attachment;filename=%s" % zip_filename}
-        )
+        if ZIP:
+            io = BytesIO()
+            zip_filename = "final_archive.zip"
+            with zipfile.ZipFile(io, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+                for file_path in [voice_bot_path, image_bot_path]:
+                    zf.write(file_path)
+                zf.close()
+            print("Total time", time.time() - total_time)
+            return StreamingResponse(
+                iter([io.getvalue()]),
+                media_type="application/x-zip-compressed",
+                headers={"Content-Disposition": f"attachment;filename=%s" % zip_filename}
+            )
+        else:
+            print("Total time", time.time() - total_time)
+            return {
+                "voice": voice_bot_path,
+                "image": image_bot_path
+            }
     else:
         return {"message": "No face detected"}