refactor: cto for better optimization

src/api/nto_api.py

@@ -57,6 +57,7 @@ class NecklaceTryOnIDEntity(BaseModel):
     necklaceImageId: str
     necklaceCategory: str
     storename: str
+
     api_token: str
 
 
@@ -128,55 +129,66 @@ 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(...)):
+                         mask: UploadFile = File(...), 
+                         clothing_type: str = Form(...)):
     logger.info("-" * 50)
     logger.info(">>> CLOTHING TRY ON STARTED <<<")
     start_time = time.time()
 
+    # Helper function to convert image to base64
+    def image_to_base64(img: Image.Image, format="WEBP", quality=85) -> str:
+        with BytesIO() as buffer:
+            img.save(buffer, format=format, quality=quality)
+            return f"data:image/{format.lower()};base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+
     try:
-        image_bytes = await image.read()
-        mask_bytes = await mask.read()
-        image, mask = Image.open(BytesIO(image_bytes)).convert("RGB"), Image.open(BytesIO(mask_bytes)).convert("RGB")
+        # Load images concurrently using asyncio
+        image_data, mask_data = await asyncio.gather(
+            image.read(),
+            mask.read()
+        )
+        
+        # Convert bytes to PIL Images
+        image = Image.open(BytesIO(image_data)).convert("RGB")
+        mask = Image.open(BytesIO(mask_data)).convert("RGB")
         logger.info(">>> IMAGES LOADED SUCCESSFULLY <<<")
     except Exception as e:
         logger.error(f">>> IMAGE LOADING ERROR: {str(e)} <<<")
-        return JSONResponse(status_code=500, content={"error": f"Error reading image or mask", "code": 500})
+        return JSONResponse(status_code=500, content={"error": "Error reading image or mask", "code": 500})
 
     try:
+        # Process images
         actual_image = image.copy()
         jewellery_mask = Image.fromarray(np.bitwise_and(np.array(mask), np.array(image)))
         arr_orig = np.array(grayscale(mask))
 
-        image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
-        image = Image.fromarray(image).resize((512, 512))
+        # Process image with inpainting
+        image = Image.fromarray(
+            cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
+        ).resize((512, 512))
 
+        # Process mask
         arr = arr_orig.copy()
         mask_y = np.where(arr == arr[arr != 0][0])[0][0]
         arr[mask_y:, :] = 255
-
         mask = Image.fromarray(arr).resize((512, 512))
+        
         logger.info(">>> IMAGE PROCESSING COMPLETED <<<")
     except Exception as e:
         logger.error(f">>> IMAGE PROCESSING ERROR: {str(e)} <<<")
-        return JSONResponse(status_code=500,
-                            content={"error": f"Error processing image or mask", "code": 500})
+        return JSONResponse(status_code=500, content={"error": "Error processing image or mask", "code": 500})
 
     try:
-        mask_img_base_64, act_img_base_64 = BytesIO(), BytesIO()
-        mask.save(mask_img_base_64, format="WEBP")
-        image.save(act_img_base_64, format="WEBP")
-        mask_bytes_ = base64.b64encode(mask_img_base_64.getvalue()).decode("utf-8")
-        image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
-
-        mask_data_uri = f"data:image/webp;base64,{mask_bytes_}"
-        image_data_uri = f"data:image/webp;base64,{image_bytes_}"
+        # Convert images to base64 more efficiently
+        mask_data_uri = image_to_base64(mask)
+        image_data_uri = image_to_base64(image)
         logger.info(">>> IMAGE ENCODING COMPLETED <<<")
     except Exception as e:
         logger.error(f">>> IMAGE ENCODING ERROR: {str(e)} <<<")
-        return JSONResponse(status_code=500,
-                            content={"error": f"Error encoding images", "code": 500})
+        return JSONResponse(status_code=500, content={"error": "Error encoding images", "code": 500})
 
-    input = {
+    # Prepare replicate input
+    input_data = {
         "mask": mask_data_uri,
         "image": image_data_uri,
         "prompt": f"Dull {clothing_type}, non-reflective clothing, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple, perfect eyes, perfect face, perfect body, high quality, realistic, photorealistic, high resolution,traditional full sleeve blouse",
@@ -185,33 +197,45 @@ async def clothing_try_on(image: UploadFile = File(...),
     }
 
     try:
-        output = replicate_run_cto(input)
+        output = replicate_run_cto(input_data)
         logger.info(">>> REPLICATE PROCESSING COMPLETED <<<")
     except Exception as e:
         logger.error(f">>> REPLICATE PROCESSING ERROR: {str(e)} <<<")
-        return JSONResponse(content={"error": f"Error running clothing try on", "code": 500}, status_code=500)
+        return JSONResponse(content={"error": "Error running clothing try on", "code": 500}, status_code=500)
 
     try:
-        response = requests.get(output[0])
-        output_image = Image.open(BytesIO(response.content)).resize(actual_image.size)
-        output_image = np.bitwise_and(np.array(output_image),
-                                      np.bitwise_not(np.array(Image.fromarray(arr_orig).convert("RGB"))))
-        result = Image.fromarray(np.bitwise_or(np.array(output_image), np.array(jewellery_mask)))
-
-        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')
+        async with aiohttp.ClientSession() as session:
+            async with session.get(output[0]) as response:
+                output_bytes = await response.read()
+        
+        output_image = Image.open(BytesIO(output_bytes)).resize(actual_image.size)
+        
+        # Process final image
+        output_array = np.bitwise_and(
+            np.array(output_image),
+            np.bitwise_not(np.array(Image.fromarray(arr_orig).convert("RGB")))
+        )
+        result = Image.fromarray(np.bitwise_or(output_array, np.array(jewellery_mask)))
+        
+        # Convert result to base64
+        result_base64 = image_to_base64(result)
+        
         total_inference_time = round((time.time() - start_time), 2)
         logger.info(">>> OUTPUT IMAGE PROCESSING COMPLETED <<<")
 
         response = {
-            "output": f"data:image/WEBP;base64,{base_64_output}",
+            "output": result_base64,
             "code": 200,
             "inference_time": total_inference_time
         }
     except Exception as e:
         logger.error(f">>> OUTPUT IMAGE PROCESSING ERROR: {str(e)} <<<")
-        return JSONResponse(status_code=500, content={"error": f"Error processing output image", "code": 500})
+        return JSONResponse(status_code=500, content={"error": "Error processing output image", "code": 500})
+    finally:
+        # Clean up resources
+        if 'output_image' in locals(): del output_image
+        if 'output_array' in locals(): del output_array
+        gc.collect()
 
     logger.info(f">>> TOTAL INFERENCE TIME: {total_inference_time}s <<<")
     logger.info(">>> REQUEST COMPLETED SUCCESSFULLY <<<")