add: nto mannequin, rt_nto

app.py

@@ -9,6 +9,7 @@ from datetime import datetime, timedelta, timezone
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse
 
 from src.api.image_prep_api import preprocessing_router
 from src.api.nto_api import nto_cto_router
@@ -35,7 +36,7 @@ async def verify_login_token(credentials: HTTPAuthorizationCredentials = Depends
         time_difference = current_time - created_at
 
         if time_difference <= timedelta(minutes=30):
-            return {"status": "Authorized", "message": "Token is valid"}
+            return JSONResponse({"status": "Authorized", "message": "Token is valid"}, status_code=200)
 
         raise HTTPException(status_code=401, detail="Unauthorized: Token expired")
 

examples/rt_nto_example.py

@@ -0,0 +1,86 @@
+"""
+project @ CTO_TCP_ZERO_GPU
+created @ 2024-11-18
+author  @ github.com/ishworrsubedii
+"""
+import asyncio
+import aiohttp
+import json
+from pathlib import Path
+
+
+async def test_rt_nto(
+        image_path: str,
+        necklace_ids: list,
+        categories: list,
+        storename: str,
+        api_url: str = "http://localhost:8000/rt_nto"
+):
+    image_content = None
+    with open(image_path, 'rb') as f:
+        image_content = f.read()
+
+    data = aiohttp.FormData()
+    data.add_field('necklace_id_list', ','.join(necklace_ids))
+    data.add_field('category_list', ','.join(categories))
+    data.add_field('storename', storename)
+    data.add_field('image',
+                   image_content,
+                   filename=Path(image_path).name,
+                   content_type='image/jpeg')
+
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.post(api_url, data=data) as response:
+                if response.status != 200:
+                    print(f"Server returned status code: {response.status}")
+                    try:
+                        error_content = await response.text()
+                        print(f"Error content: {error_content}")
+                    except Exception as e:
+                        print(f"Could not read error content: {e}")
+                    return
+
+                async for line in response.content:
+                    if line:
+                        try:
+                            result = json.loads(line.decode('utf-8'))
+                            if result.get('code') == 200:
+                                print(f"Success - Progress: {result['progress']}")
+                                print(f"Necklace ID: {result['necklace_id']}")
+                                print(f"Category: {result['category']}")
+                                print(f"Output URL: {result['output']}")
+                                print(f"Mask URL: {result['mask']}")
+                                print(f"Inference Time: {result['inference_time']}s")
+                            else:
+                                print(f"Error - Progress: {result.get('progress')}")
+                                print(f"Error Details: {result.get('error')}")
+                            print("-" * 50)
+                        except json.JSONDecodeError as e:
+                            print(f"Error decoding response: {e}")
+                            print(f"Raw line: {line}")
+        except aiohttp.ClientError as e:
+            print(f"Connection error: {e}")
+        except Exception as e:
+            print(f"Unexpected error: {e}")
+
+
+async def main():
+    image_path = "/home/ishwor/Downloads/download (7).png"
+    necklace_ids = ["ST01SHO001", "ST01SHO002", "ST01SHO003"]
+    categories = ["Short Necklaces", "Short Necklaces", "Short Necklaces"]
+    storename = "Store1"
+
+    try:
+        await test_rt_nto(
+            image_path=image_path,
+            necklace_ids=necklace_ids,
+            categories=categories,
+            storename=storename
+        )
+    except Exception as e:
+        print(f"Error in main: {e}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/api/batch_api.py

@@ -4,6 +4,7 @@ created @ 2024-11-14
 author  @ github.com/ishworrsubedii
 """
 import base64
+import gc
 import time
 from io import BytesIO
 import json
@@ -15,8 +16,12 @@ from fastapi.routing import APIRouter
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import List
+
+from fastapi.responses import JSONResponse
+
+from src.utils import returnBytesData
 from src.utils.logger import logger
-from src.api.nto_api import pipeline, replicate_run_cto
+from src.api.nto_api import pipeline, replicate_run_cto, supabase_upload_and_return_url
 
 batch_router = APIRouter()
 
@@ -36,7 +41,7 @@ async def rt_cto(
     try:
         clothing_list = [item.strip() for item in c_list.split(",")]
         logger.info(f">>> CLOTHING LIST: {clothing_list} <<<")
-        
+
         image_bytes = await image.read()
         pil_image = Image.open(BytesIO(image_bytes)).convert("RGB")
         logger.info(">>> IMAGE LOADED SUCCESSFULLY <<<")
@@ -77,7 +82,7 @@ async def rt_cto(
         for idx, clothing_type in enumerate(clothing_list):
             if not clothing_type:
                 continue
-                
+
             input = {
                 "mask": mask_data_uri,
                 "image": image_data_uri,
@@ -89,10 +94,10 @@ async def rt_cto(
             try:
                 output = replicate_run_cto(input)
                 logger.info(f">>> REPLICATE PROCESSING COMPLETED FOR {clothing_type} <<<")
-                
+
                 output_url = str(output[0]) if output and output[0] else None
                 total_inference_time = round((time.time() - start_time), 2)
-                
+
                 result = {
                     "code": 200,
                     "output": output_url,
@@ -102,7 +107,7 @@ async def rt_cto(
                 }
                 yield json.dumps(result) + "\n"
                 await asyncio.sleep(0.1)
-                
+
             except Exception as e:
                 logger.error(f">>> REPLICATE PROCESSING ERROR: {str(e)} <<<")
                 error_result = {
@@ -125,3 +130,114 @@ async def rt_cto(
             "Transfer-Encoding": "chunked"
         }
     )
+
+
+@batch_router.post("/rt_nto")
+async def rt_nto(
+        image: UploadFile = File(...),
+        necklace_id_list: str = Form(...),
+        category_list: str = Form(...),
+        storename: str = Form(...)
+):
+    logger.info("-" * 50)
+    logger.info(">>> REAL-TIME NECKLACE TRY ON STARTED <<<")
+    logger.info(f"Parameters: storename={storename}, categories={category_list}, necklace_ids={necklace_id_list}")
+
+    try:
+        # Parse the lists
+        necklace_ids = [id.strip() for id in necklace_id_list.split(",")]
+        categories = [cat.strip() for cat in category_list.split(",")]
+        if len(necklace_ids) != len(categories):
+            return JSONResponse(
+                content={"error": "Number of necklace IDs must match number of categories", "code": 400},
+                status_code=400
+            )
+
+        # Load the source image
+        image_bytes = await image.read()
+        source_image = Image.open(BytesIO(image_bytes))
+        logger.info(">>> SOURCE IMAGE LOADED SUCCESSFULLY <<<")
+    except Exception as e:
+        logger.error(f">>> INITIAL SETUP ERROR: {str(e)} <<<")
+        return JSONResponse(
+            content={"error": "Error in initial setup", "details": str(e), "code": 500},
+            status_code=500
+        )
+
+    async def generate():
+        start_time = time.time()
+
+        for idx, (necklace_id, category) in enumerate(zip(necklace_ids, categories)):
+            try:
+                jewellery_url = f"https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/Stores/{storename}/{category}/image/{necklace_id}.png"
+                jewellery = Image.open(returnBytesData(url=jewellery_url))
+                logger.info(f">>> JEWELLERY IMAGE {necklace_id} LOADED SUCCESSFULLY <<<")
+
+                # Process the necklace try-on
+                result, headetText, mask = await pipeline.necklaceTryOn_(
+                    image=source_image,
+                    jewellery=jewellery,
+                    storename=storename
+                )
+
+                if result is None:
+                    error_result = {
+                        "error": "No face detected in the image",
+                        "code": 400,
+                        "necklace_id": necklace_id,
+                        "category": category,
+                        "progress": f"{idx + 1}/{len(necklace_ids)}"
+                    }
+                    yield json.dumps(error_result) + "\n"
+                    continue
+
+                # Upload results concurrently
+                logger.info(">>> UPLOADING RESULTS <<<")
+                upload_tasks = [
+                    supabase_upload_and_return_url(prefix="necklace_try_on", image=result),
+                    supabase_upload_and_return_url(prefix="necklace_try_on_mask", image=mask)
+                ]
+                result_url, mask_url = await asyncio.gather(*upload_tasks)
+
+                total_inference_time = round((time.time() - start_time), 2)
+                logger.info(f">>> UPLOADING COMPLETED FOR {necklace_id} <<<")
+
+                result = {
+                    "code": 200,
+                    "output": result_url,
+                    "mask": mask_url,
+                    "inference_time": total_inference_time,
+                    "necklace_id": necklace_id,
+                    "category": category,
+                    "progress": f"{idx + 1}/{len(necklace_ids)}"
+                }
+                yield json.dumps(result) + "\n"
+                await asyncio.sleep(0.1)
+
+                del result
+                del mask
+                gc.collect()
+
+            except Exception as e:
+                logger.error(f">>> PROCESSING ERROR FOR {necklace_id}: {str(e)} <<<")
+                error_result = {
+                    "error": f"Error processing necklace {necklace_id}",
+                    "details": str(e),
+                    "code": 500,
+                    "necklace_id": necklace_id,
+                    "category": category,
+                    "progress": f"{idx + 1}/{len(necklace_ids)}"
+                }
+                yield json.dumps(error_result) + "\n"
+                await asyncio.sleep(0.1)
+
+    return StreamingResponse(
+        generate(),
+        media_type="application/x-ndjson",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+            "Transfer-Encoding": "chunked"
+        }
+    )

src/api/nto_api.py

@@ -749,7 +749,7 @@ async def mannequin_nto(necklace_try_on_id: NecklaceTryOnIDEntity = Depends(pars
             "code": 404}, status_code=404)
 
     try:
-        result = await pipeline.necklaceTryOnMannequin_(image=image, jewellery=jewellery)
+        result,resized_img = await pipeline.necklaceTryOnMannequin_(image=image, jewellery=jewellery)
 
         if result is None:
             logger.error(">>> NO FACE DETECTED IN THE IMAGE <<<")

src/components/necklaceTryOn.py

@@ -107,6 +107,7 @@ class NecklaceTryOn:
             upper_half = image_np[:middle_point, :]
             lower_half = image_np[middle_point:, :]
 
+            upper_half = cv2.resize(upper_half, (upper_half.shape[1] * 3, upper_half.shape[0] * 3))
             copy_upper = upper_half.copy()
 
             # Apply pose detection on upper half
@@ -145,10 +146,12 @@ class NecklaceTryOn:
 
             result_upper = cvzone.overlayPNG(copy_upper, jewellery, (avg_x1, y_coordinate))
 
-            final_result = np.vstack((result_upper, lower_half))
+            final_result = cv2.resize(result_upper, (image_np.shape[1], image_np.shape[0] - middle_point))
+
+            final_result = np.vstack((final_result, lower_half))
 
             gc.collect()
-            return Image.fromarray(final_result.astype(np.uint8))
+            return Image.fromarray(final_result.astype(np.uint8)), Image.fromarray(result_upper.astype(np.uint8))
 
         except Exception as e:
             print(f"Error: {e}")

src/pipelines/completePipeline.py

@@ -30,7 +30,7 @@ class Pipeline:
                                                                                  storename=storename)
         return [result, headerText, mask]
 
-    async def necklaceTryOnMannequin_(self, image: Image.Image, jewellery: Image.Image):
-        result = self.necklaceTryOnObj.necklaceTryOnMannequin(image, jewellery)
+    def necklaceTryOnMannequin_(self, image: Image.Image, jewellery: Image.Image):
+        result, resized_image = self.necklaceTryOnObj.necklaceTryOnMannequin(image, jewellery)
 
-        return result
+        return result, resized_image