add- fit to image

runway.py

@@ -24,7 +24,7 @@ async def root():
     """
     Root endpoint for basic health check.
     """
-    return {"message": "InstructPix2Pix API is running. Use POST /inpaint/ or /inpaint-with-reference/ to edit images."}
+    return {"message": "InstructPix2Pix API is running. Use POST /inpaint/, /inpaint-with-reference/, or /fit-image-to-mask/ to edit images."}
 
 def prepare_guided_image(original_image: Image, reference_image: Image, mask_image: Image) -> Image:
     """
@@ -40,30 +40,16 @@ def prepare_guided_image(original_image: Image, reference_image: Image, mask_ima
     Returns:
         Image: The blended image to guide inpainting.
     """
-    # Convert images to numpy arrays
     original_array = np.array(original_image)
     reference_array = np.array(reference_image)
-    mask_array = np.array(mask_image) / 255.0  # Normalize to [0, 1]
-
-    # Expand mask to RGB channels
+    mask_array = np.array(mask_image) / 255.0
     mask_array = mask_array[:, :, np.newaxis]
-
-    # Blend: use original where mask=0 (black), reference where mask=1 (white)
     blended_array = original_array * (1 - mask_array) + reference_array * mask_array
-    blended_array = blended_array.astype(np.uint8)
-
-    return Image.fromarray(blended_array)
+    return Image.fromarray(blended_array.astype(np.uint8))
 
 def soften_mask(mask_image: Image, softness: int = 5) -> Image:
     """
     Soften the edges of the mask for smoother transitions.
-    
-    Args:
-        mask_image (Image): The original mask (grayscale, L mode).
-        softness (int): Size of the Gaussian blur kernel for softening edges.
-    
-    Returns:
-        Image: The softened mask.
     """
     from PIL import ImageFilter
     return mask_image.filter(ImageFilter.GaussianBlur(radius=softness))
@@ -72,24 +58,59 @@ def generate_rectangular_mask(image_size: tuple, x1: int = 100, y1: int = 100, x
     """
     Generate a rectangular mask matching the image dimensions.
     - Black (0) for areas to keep, white (255) for areas to inpaint.
-    
-    Args:
-        image_size (tuple): Tuple of (width, height) of the original image.
-        x1, y1 (int): Top-left corner coordinates of the rectangle.
-        x2, y2 (int): Bottom-right corner coordinates of the rectangle.
-    
-    Returns:
-        Image: The generated mask in grayscale (L mode).
     """
-    # Create a blank black mask (0 = keep)
     mask = Image.new("L", image_size, 0)
     draw = ImageDraw.Draw(mask)
-    
-    # Draw a white rectangle (255 = inpaint)
     draw.rectangle([x1, y1, x2, y2], fill=255)
-    
     return mask
 
+def fit_image_to_mask(original_image: Image, reference_image: Image, mask_x1: int, mask_y1: int, mask_x2: int, mask_y2: int) -> tuple:
+    """
+    Fit the reference image into the masked region of the original image.
+    
+    Args:
+        original_image (Image): The original image (RGB).
+        reference_image (Image): The image to fit into the masked region (RGB).
+        mask_x1, mask_y1, mask_x2, mask_y2 (int): Coordinates of the masked region.
+    
+    Returns:
+        tuple: (guided_image, mask_image) - The image with the fitted reference and the corresponding mask.
+    """
+    # Calculate mask dimensions
+    mask_width = mask_x2 - mask_x1
+    mask_height = mask_y2 - mask_y1
+
+    # Resize reference image to fit the mask while preserving aspect ratio
+    ref_width, ref_height = reference_image.size
+    aspect_ratio = ref_width / ref_height
+
+    if mask_width / mask_height > aspect_ratio:
+        # Fit to height
+        new_height = mask_height
+        new_width = int(new_height * aspect_ratio)
+    else:
+        # Fit to width
+        new_width = mask_width
+        new_height = int(new_width / aspect_ratio)
+
+    # Resize reference image
+    reference_image_resized = reference_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    # Create a copy of the original image to paste the reference image onto
+    guided_image = original_image.copy()
+
+    # Calculate position to center the resized image in the mask
+    paste_x = mask_x1 + (mask_width - new_width) // 2
+    paste_y = mask_y1 + (mask_height - new_height) // 2
+
+    # Paste the resized reference image onto the original image
+    guided_image.paste(reference_image_resized, (paste_x, paste_y))
+
+    # Generate the mask for inpainting (white in the pasted region)
+    mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+
+    return guided_image, mask_image
+
 @app.post("/inpaint/")
 async def inpaint_image(
     image: UploadFile = File(...),
@@ -101,36 +122,20 @@ async def inpaint_image(
 ):
     """
     Endpoint for image inpainting using a text prompt and autogenerated mask.
-    - `image`: Original image file (PNG/JPG).
-    - `prompt`: Text prompt describing the desired output.
-    - `mask_x1, mask_y1, mask_x2, mask_y2`: Coordinates for the rectangular mask (default: 100,100 to 200,200).
-    
-    Returns:
-    - The inpainted image as a PNG file.
     """
     try:
-        # Load the uploaded image
         image_bytes = await image.read()
         original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-        # Generate the mask based on image dimensions and provided coordinates
         mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
-
-        # Perform inpainting using the pipeline
         result = pipe(prompt=prompt, image=original_image, mask_image=mask_image).images[0]
-
-        # Convert result to bytes for response
         result_bytes = io.BytesIO()
         result.save(result_bytes, format="PNG")
         result_bytes.seek(0)
-
-        # Return the image as a streaming response
         return StreamingResponse(
             result_bytes,
             media_type="image/png",
             headers={"Content-Disposition": "attachment; filename=inpainted_image.png"}
         )
-    
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error during inpainting: {e}")
 
@@ -145,9 +150,53 @@ async def inpaint_with_reference(
     mask_y2: int = 200
 ):
     """
-    Endpoint for replacing masked areas with reference image content, refined to look natural, using an autogenerated mask.
-    - `image`: Original image file (PNG/JPG).
-    - `reference_image`: Reference image to guide the replacement (PNG/JPG).
+    Endpoint for replacing masked areas with reference image content, refined to look natural.
+    """
+    try:
+        image_bytes = await image.read()
+        reference_bytes = await reference_image.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
+
+        if original_image.size != reference_image.size:
+            reference_image = reference_image.resize(original_image.size, Image.Resampling.LANCZOS)
+
+        mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+        softened_mask = soften_mask(mask_image, softness=5)
+        guided_image = prepare_guided_image(original_image, reference_image, softened_mask)
+        result = pipe(
+            prompt=prompt,
+            image=guided_image,
+            mask_image=softened_mask,
+            strength=0.75,
+            guidance_scale=7.5
+        ).images[0]
+
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=natural_inpaint_image.png"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during natural inpainting: {e}")
+
+@app.post("/fit-image-to-mask/")
+async def fit_image_to_mask(
+    image: UploadFile = File(...),
+    reference_image: UploadFile = File(...),
+    prompt: str = "Blend the fitted image naturally into the scene, matching style and lighting.",
+    mask_x1: int = 100,
+    mask_y1: int = 100,
+    mask_x2: int = 200,
+    mask_y2: int = 200
+):
+    """
+    Endpoint for fitting a reference image into a masked region of the original image, refined to look natural.
+    - `image`: Original image file (PNG/JPG), e.g., a table.
+    - `reference_image`: Image to fit into the masked region (PNG/JPG), e.g., a cat.
     - `prompt`: Text prompt for inpainting refinement.
     - `mask_x1, mask_y1, mask_x2, mask_y2`: Coordinates for the rectangular mask (default: 100,100 to 200,200).
     
@@ -155,49 +204,38 @@ async def inpaint_with_reference(
     - The resulting image as a PNG file.
     """
     try:
-        # Load the uploaded image and reference image
+        # Load the uploaded images
         image_bytes = await image.read()
         reference_bytes = await reference_image.read()
-
         original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
 
-        # Ensure reference image matches original image dimensions
-        if original_image.size != reference_image.size:
-            reference_image = reference_image.resize(original_image.size, Image.Resampling.LANCZOS)
-
-        # Generate the mask based on image dimensions and provided coordinates
-        mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+        # Fit the reference image into the masked region
+        guided_image, mask_image = fit_image_to_mask(original_image, reference_image, mask_x1, mask_y1, mask_x2, mask_y2)
 
         # Soften the mask for smoother transitions
         softened_mask = soften_mask(mask_image, softness=5)
 
-        # Prepare the initial guided image by blending reference content into the masked area
-        guided_image = prepare_guided_image(original_image, reference_image, softened_mask)
-
-        # Perform inpainting to refine the result and make it look natural
+        # Perform inpainting to blend the fitted image naturally
         result = pipe(
             prompt=prompt,
             image=guided_image,
-            mask_image=softened_mask,  # Use softened mask for inpainting
-            strength=0.75,  # Control how much inpainting modifies the image (0.0 to 1.0)
-            guidance_scale=7.5  # Control how closely the prompt is followed
+            mask_image=softened_mask,
+            strength=0.75,
+            guidance_scale=7.5
         ).images[0]
 
         # Convert result to bytes for response
         result_bytes = io.BytesIO()
         result.save(result_bytes, format="PNG")
         result_bytes.seek(0)
-
-        # Return the image as a streaming response
         return StreamingResponse(
             result_bytes,
             media_type="image/png",
-            headers={"Content-Disposition": "attachment; filename=natural_inpaint_image.png"}
+            headers={"Content-Disposition": "attachment; filename=fitted_image.png"}
         )
-    
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error during natural inpainting: {e}")
+        raise HTTPException(status_code=500, detail=f"Error during fitting and inpainting: {e}")
 
 if __name__ == "__main__":
     import uvicorn