add: inference time

src/api/nto_api.py

@@ -3,6 +3,7 @@ project @ NTO-TCP-HF
 created @ 2024-10-28
 author  @ github/ishworrsubedii
 """
+import time
 import cv2
 import numpy as np
 from PIL.ImageOps import grayscale
@@ -53,6 +54,7 @@ class NecklaceTryOnIDEntity(BaseModel):
 
 @nto_cto_router.post("/clothingTryOnV2")
 async def clothing_try_on_v2(image: UploadFile = File(...), clothing_type: str = Form(...)):
+    start_time = time.time()
     image_bytes = await image.read()
     image = Image.open(BytesIO(image_bytes)).convert("RGB")
 
@@ -76,10 +78,12 @@ async def clothing_try_on_v2(image: UploadFile = File(...), clothing_type: str =
     }
 
     output = replicate_run_cto(input)
+    total_inference_time = round((time.time() - start_time), 2)
 
     response = {
         "output": f"{output[0]}",
-        'code': 200
+        'code': 200,
+        "inference_time": total_inference_time
     }
 
     return JSONResponse(content=response, status_code=200)
@@ -88,6 +92,7 @@ async def clothing_try_on_v2(image: UploadFile = File(...), clothing_type: str =
 @nto_cto_router.post("/clothingTryOn")
 async def clothing_try_on(image: UploadFile = File(...),
                           mask: UploadFile = File(...), clothing_type: str = Form(...)):
+    start_time = time.time()
     image_bytes = await image.read()
     mask_bytes = await mask.read()
     image, mask = Image.open(BytesIO(image_bytes)).convert("RGB"), Image.open(
@@ -142,10 +147,12 @@ async def clothing_try_on(image: UploadFile = File(...),
     in_mem_file = BytesIO()
     result.save(in_mem_file, format="WEBP", quality=85)
     base_64_output = base64.b64encode(in_mem_file.getvalue()).decode('utf-8')
+    total_inference_time = round((time.time() - start_time), 2)
 
     response = {
         "output": f"data:image/WEBP;base64,{base_64_output}",
-        'code': 200
+        "code": 200,
+        "inference_time": total_inference_time
     }
 
     return JSONResponse(content=response, status_code=200)
@@ -297,6 +304,7 @@ async def parse_necklace_try_on_id(necklaceImageId: str = Form(...),
 @nto_cto_router.post("/necklaceTryOnID")
 async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends(parse_necklace_try_on_id),
                              image: UploadFile = File(...)):
+    start_time = time.time()
     data, _ = supabase.table("APIKeyList").select("*").filter("API_KEY", "eq",
                                                               necklace_try_on_id.api_token).execute()
 
@@ -328,11 +336,14 @@ async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends
         mask.save(inMemFileMask, format="WEBP", quality=85)
         outputBytes = inMemFile.getvalue()
         maskBytes = inMemFileMask.getvalue()
+        creditResponse = deductAndTrackCredit(storename=necklace_try_on_id.storename, endpoint="/necklaceTryOnID")
+        total_backend_time = round((time.time() - start_time), 2)
         response = {
             "output": f"data:image/WEBP;base64,{base64.b64encode(outputBytes).decode('utf-8')}",
-            "mask": f"data:image/WEBP;base64,{base64.b64encode(maskBytes).decode('utf-8')}"
+            "mask": f"data:image/WEBP;base64,{base64.b64encode(maskBytes).decode('utf-8')}",
+            "inference_time": total_backend_time
         }
-        creditResponse = deductAndTrackCredit(storename=necklace_try_on_id.storename, endpoint="/necklaceTryOnID")
+
         if creditResponse == "No Credits Available":
             response = {
                 "error": "No Credits Remaining"
@@ -384,6 +395,7 @@ async def necklace_try_on_with_points(necklace_try_on_id: NecklaceTryOnIDEntity
                                       left_y: int = Form(...),
                                       right_x: int = Form(...),
                                       right_y: int = Form(...)):
+    start_time = time.time()
     imageBytes = await image.read()
 
     jewellery_url = f"https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/Stores/{necklace_try_on_id.storename}/{necklace_try_on_id.necklaceCategory}/image/{necklace_try_on_id.necklaceImageId}.png"
@@ -409,12 +421,14 @@ async def necklace_try_on_with_points(necklace_try_on_id: NecklaceTryOnIDEntity
     mask.save(inMemFileMask, format="WEBP", quality=85)
     outputBytes = inMemFile.getvalue()
     maskBytes = inMemFileMask.getvalue()
+
+    creditResponse = deductAndTrackCredit(storename=necklace_try_on_id.storename, endpoint="/necklaceTryOnID")
+    total_inference_time = round((time.time() - start_time), 2)
     response = {
         "output": f"data:image/WEBP;base64,{base64.b64encode(outputBytes).decode('utf-8')}",
-        "mask": f"data:image/WEBP;base64,{base64.b64encode(maskBytes).decode('utf-8')}"
+        "mask": f"data:image/WEBP;base64,{base64.b64encode(maskBytes).decode('utf-8')}",
+        "inference_time": total_inference_time
     }
-
-    creditResponse = deductAndTrackCredit(storename=necklace_try_on_id.storename, endpoint="/necklaceTryOnID")
     if creditResponse == "No Credits Available":
         response = {
             "error": "No Credits Remaining"