improve test, add extract to api

client_api/api_utils.py

@@ -8,6 +8,7 @@ import base64, io
 from io import BytesIO
 from typing import List, Tuple, Optional
 import numpy as np
+import requests
 
 
 class InpaintingWhen(Enum):
@@ -151,7 +152,7 @@ class FaceSwapRequest(BaseModel):
 
 class FaceSwapResponse(BaseModel):
     images: List[str] = Field(description="base64 swapped image", default=None)
-    infos: List[str]
+    infos: Optional[List[str]]  # not really used atm
 
     @property
     def pil_images(self) -> Image.Image:
@@ -171,6 +172,23 @@ class FaceSwapCompareRequest(BaseModel):
     )
 
 
+class FaceSwapExtractRequest(BaseModel):
+    images: List[str] = Field(
+        description="base64 reference image",
+        examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."],
+        default=None,
+    )
+    postprocessing: Optional[PostProcessingOptions]
+
+
+class FaceSwapExtractResponse(BaseModel):
+    images: List[str] = Field(description="base64 face images", default=None)
+
+    @property
+    def pil_images(self) -> Image.Image:
+        return [base64_to_pil(img) for img in self.images]
+
+
 def pil_to_base64(img: Image.Image) -> np.array:  # type:ignore
     if isinstance(img, str):
         img = Image.open(img)
@@ -192,3 +210,20 @@ def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]:
         # if no data URL scheme, just decode
         img_bytes = base64.b64decode(base64str)
     return Image.open(io.BytesIO(img_bytes))
+
+
+def compare_faces(
+    image1: Image.Image, image2: Image.Image, base_url: str = "http://localhost:7860"
+) -> float:
+    request = FaceSwapCompareRequest(
+        image1=pil_to_base64(image1),
+        image2=pil_to_base64(image2),
+    )
+
+    result = requests.post(
+        url=f"{base_url}/faceswaplab/compare",
+        data=request.json(),
+        headers={"Content-Type": "application/json; charset=utf-8"},
+    )
+
+    return float(result.text)

client_api/faceswaplab_api_example.py

@@ -7,10 +7,16 @@ from api_utils import (
     pil_to_base64,
     InpaintingWhen,
     FaceSwapCompareRequest,
+    FaceSwapExtractRequest,
+    FaceSwapExtractResponse,
 )
 
 address = "http://127.0.0.1:7860"
 
+
+#############################
+# FaceSwap
+
 # First face unit :
 unit1 = FaceSwapUnit(
     source_img=pil_to_base64("../references/man.png"),  # The face you want to use
@@ -41,7 +47,7 @@ request = FaceSwapRequest(
     image=pil_to_base64("test_image.png"), units=[unit1, unit2], postprocessing=pp
 )
 
-
+# Face Swap
 result = requests.post(
     url=f"{address}/faceswaplab/swap_face",
     data=request.json(),
@@ -52,6 +58,8 @@ response = FaceSwapResponse.parse_obj(result.json())
 for img in response.pil_images:
     img.show()
 
+#############################
+# Comparison
 
 request = FaceSwapCompareRequest(
     image1=pil_to_base64("../references/man.png"),
@@ -65,3 +73,21 @@ result = requests.post(
 )
 
 print("similarity", result.text)
+
+#############################
+# Extraction
+
+# Prepare the request
+request = FaceSwapExtractRequest(
+    images=[pil_to_base64(response.pil_images[0])], postprocessing=pp
+)
+
+result = requests.post(
+    url=f"{address}/faceswaplab/extract",
+    data=request.json(),
+    headers={"Content-Type": "application/json; charset=utf-8"},
+)
+response = FaceSwapExtractResponse.parse_obj(result.json())
+
+for img in response.pil_images:
+    img.show()