refactor: cto for better optimization

src/api/nto_api.py

@@ -148,7 +148,6 @@ async def clothing_try_on(image: UploadFile = File(...),
             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 <<<")
@@ -157,17 +156,14 @@ async def clothing_try_on(image: UploadFile = File(...),
         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))
 
-        # 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
@@ -179,7 +175,6 @@ async def clothing_try_on(image: UploadFile = File(...),
         return JSONResponse(status_code=500, content={"error": "Error processing image or mask", "code": 500})
 
     try:
-        # 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 <<<")
@@ -187,7 +182,6 @@ async def clothing_try_on(image: UploadFile = File(...),
         logger.error(f">>> IMAGE ENCODING ERROR: {str(e)} <<<")
         return JSONResponse(status_code=500, content={"error": "Error encoding images", "code": 500})
 
-    # Prepare replicate input
     input_data = {
         "mask": mask_data_uri,
         "image": image_data_uri,
@@ -204,18 +198,30 @@ async def clothing_try_on(image: UploadFile = File(...),
         return JSONResponse(content={"error": "Error running clothing try on", "code": 500}, status_code=500)
 
     try:
+        # Fetch the output image URL
+        output_url = output[0]
+        if not isinstance(output_url, str):
+            raise ValueError("Invalid output URL format from replicate")
+
+        # Download the output image
         async with aiohttp.ClientSession() as session:
-            async with session.get(output[0]) as response:
+            async with session.get(output_url) as response:
+                if response.status != 200:
+                    raise HTTPException(status_code=response.status, 
+                                     detail="Failed to fetch output image")
                 output_bytes = await response.read()
         
+        # Process the output image
         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 arrays and process final image
+        mask_array = np.array(Image.fromarray(arr_orig).convert("RGB"))
+        output_array = np.array(output_image)
+        
+        # Perform bitwise operations
+        mask_inverse = np.bitwise_not(mask_array)
+        intermediate = np.bitwise_and(output_array, mask_inverse)
+        result = Image.fromarray(np.bitwise_or(intermediate, np.array(jewellery_mask)))
         
         # Convert result to base64
         result_base64 = image_to_base64(result)
@@ -228,13 +234,21 @@ async def clothing_try_on(image: UploadFile = File(...),
             "code": 200,
             "inference_time": total_inference_time
         }
+    except ValueError as ve:
+        logger.error(f">>> OUTPUT IMAGE PROCESSING ERROR: {str(ve)} <<<")
+        return JSONResponse(status_code=500, 
+                          content={"error": "Invalid response from image generation service", 
+                                 "code": 500})
     except Exception as e:
         logger.error(f">>> OUTPUT IMAGE PROCESSING ERROR: {str(e)} <<<")
-        return JSONResponse(status_code=500, content={"error": "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
+        for var in ['output_image', 'output_array', 'intermediate', 'mask_array', 'mask_inverse']:
+            if var in locals():
+                del locals()[var]
         gc.collect()
 
     logger.info(f">>> TOTAL INFERENCE TIME: {total_inference_time}s <<<")