shelter / src /image_composite.py
CalamityJon's picture
Upload 32 files
0453188 verified
Raw
History Blame Contribute Delete
3.98 kB
"""
Image compositing and enhancement module.
"""
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np
def composite_image(foreground, background, target_size=(1200, 1600)):
"""
Composite the animal (with removed background) onto a new background.
Args:
foreground: PIL Image (RGBA) of animal with transparent background
background: PIL Image of the background scene
target_size: tuple of (width, height) for final output
Returns:
PIL Image: Composited and enhanced image
"""
# Resize background to target size
background = background.convert('RGB')
background = background.resize(target_size, Image.Resampling.LANCZOS)
# Calculate foreground size (keep animal prominent but not too large)
# Animal should take up about 60-70% of the height
fg_max_height = int(target_size[1] * 0.7)
fg_max_width = int(target_size[0] * 0.8)
# Resize foreground maintaining aspect ratio
fg_ratio = foreground.size[0] / foreground.size[1]
if foreground.size[1] > fg_max_height:
new_height = fg_max_height
new_width = int(new_height * fg_ratio)
else:
new_height = foreground.size[1]
new_width = foreground.size[0]
# Ensure width doesn't exceed max
if new_width > fg_max_width:
new_width = fg_max_width
new_height = int(new_width / fg_ratio)
foreground_resized = foreground.resize((new_width, new_height), Image.Resampling.LANCZOS)
# Calculate position to place the animal at the bottom (grounded, no margin)
x_pos = (target_size[0] - new_width) // 2 # Center horizontally
y_pos = target_size[1] - new_height # Bottom edge, no margin
# Ensure position is within bounds
y_pos = max(0, min(y_pos, target_size[1] - new_height))
# Create a copy of background to paste onto
result = background.copy()
# Paste foreground using alpha channel as mask
if foreground_resized.mode == 'RGBA':
result.paste(foreground_resized, (x_pos, y_pos), foreground_resized)
else:
# If no alpha channel, just paste normally
result.paste(foreground_resized, (x_pos, y_pos))
# Apply enhancements
result = enhance_image(result)
return result
def enhance_image(image):
"""
Apply subtle enhancements to make the image more appealing.
Args:
image: PIL Image (RGB)
Returns:
PIL Image: Enhanced image
"""
# Increase brightness slightly (5%)
enhancer = ImageEnhance.Brightness(image)
image = enhancer.enhance(1.05)
# Increase contrast slightly (10%)
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(1.10)
# Increase color saturation slightly (8%)
enhancer = ImageEnhance.Color(image)
image = enhancer.enhance(1.08)
# Apply subtle sharpening
image = image.filter(ImageFilter.UnsharpMask(radius=1, percent=50, threshold=3))
return image
def blend_edges(foreground, background, blur_radius=5):
"""
Optional: Blend the edges of the foreground with the background for smoother compositing.
Args:
foreground: PIL Image (RGBA)
background: PIL Image (RGB)
blur_radius: int, radius for edge blurring
Returns:
PIL Image: Composited image with blended edges
"""
if foreground.mode != 'RGBA':
return foreground
# Extract alpha channel
alpha = foreground.split()[3]
# Apply slight blur to alpha for softer edges
alpha_blurred = alpha.filter(ImageFilter.GaussianBlur(blur_radius))
# Recreate foreground with blurred alpha
fg_with_blur = Image.merge('RGBA', foreground.split()[:3] + (alpha_blurred,))
return fg_with_blur