Update app/api/image_router.py

app/api/image_router.py

@@ -122,37 +122,46 @@ async def update_image(
 # --- New Virtual Try-On Endpoint with File Uploads ---
 @router.post("/virtual-try-on")
 async def virtual_try_on(
-    dress_image: UploadFile = File(..., description="The dress image for try-on (PNG, JPG)"),
-    person_image: UploadFile = File(..., description="The person image for try-on (PNG, JPG)")
-) -> Response: # Return type is Response because it can be JSON or Image
+    dress_image: UploadFile = File(..., description="The dress image (Input 1)"),
+    person_image: UploadFile = File(..., description="The person image (Input 2)"),
+    shoes_image: UploadFile = File(None, description="The shoes image (Input 3) - Optional")
+) -> Response: 
     """
-    Performs a virtual try-on using an uploaded dress image and a person image.
-    Returns the try-on image directly or a JSON response with a summary if no image.
+    Performs a virtual try-on. 
+    Accepts: Dress, Person, and optionally Shoes.
     """
     logger.info("Received request to /virtual-try-on with image uploads.")
 
+    # Validate mandatory images
     if not dress_image.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="Dress image must be an image type (e.g., image/png, image/jpeg).")
+        raise HTTPException(status_code=400, detail="Dress image must be an image type.")
     if not person_image.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="Person image must be an image type (e.g., image/png, image/jpeg).")
+        raise HTTPException(status_code=400, detail="Person image must be an image type.")
+
+    # Validate shoes only if provided
+    if shoes_image and not shoes_image.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Shoes image must be an image type.")
 
     try:
         dress_image_bytes = await dress_image.read()
         person_image_bytes = await person_image.read()
+        
+        shoes_image_bytes = None
+        if shoes_image:
+            shoes_image_bytes = await shoes_image.read()
 
         # Service returns JSON summary (dict) and BytesIO object
         summary, try_on_image_bytes_io = image_service.virtual_try_on(
             dress_image_bytes=dress_image_bytes,
-            person_image_bytes=person_image_bytes
+            person_image_bytes=person_image_bytes,
+            shoes_image_bytes=shoes_image_bytes
         )
         
         if try_on_image_bytes_io:
             logger.info("Virtual try-on successful. Streaming image response.")
-            # If an image is returned, stream it directly
             return create_image_streaming_response(try_on_image_bytes_io)
         else:
             logger.warning("Virtual try-on returned no image. Sending JSON summary.")
-            # If no image, return the JSON summary as a JSON response
             return JSONResponse(content=summary)
             
     except HTTPException:
@@ -160,7 +169,6 @@ async def virtual_try_on(
     except Exception as e:
         logger.error(f"Error in /virtual-try-on: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))
-    
 
 @router.post("/generate-shoe", response_class=StreamingResponse)
 async def generate_shoe(request: ShoeGenerateRequest):