257 lines
11 KiB
Python
257 lines
11 KiB
Python
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
from dataclasses import dataclass
|
|
from typing import Tuple, Union
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@dataclass
|
|
class MaskProcessingConfig:
|
|
"""Configuration for mask processing parameters"""
|
|
mask_expansion: int = 0
|
|
edge_detection: bool = False
|
|
edge_thickness: int = 1
|
|
edge_color: Tuple[int, int, int, int] = (0, 0, 0, 255)
|
|
mask_blur: int = 5
|
|
threshold: float = 0.5
|
|
invert_generated_mask: bool = False
|
|
remove_small_regions: bool = False
|
|
small_region_size: int = 100
|
|
edge_refinement: bool = False
|
|
|
|
class EnhancedMaskProcessor:
|
|
"""Enhanced mask processing with advanced refinement techniques"""
|
|
|
|
@staticmethod
|
|
def refine_mask(mask: Image.Image, config: MaskProcessingConfig, original_image: Image.Image) -> Image.Image:
|
|
"""Enhanced mask refinement with improved edge detection and color control"""
|
|
try:
|
|
# Convert mask to numpy array
|
|
mask_np = np.array(mask)
|
|
if len(mask_np.shape) > 2:
|
|
if mask_np.shape[2] == 4:
|
|
mask_np = mask_np[:, :, 3]
|
|
else:
|
|
mask_np = mask_np[:, :, 0]
|
|
|
|
# Initial binary threshold
|
|
_, binary_mask = cv2.threshold(
|
|
mask_np,
|
|
127,
|
|
255,
|
|
cv2.THRESH_BINARY
|
|
)
|
|
|
|
# Enhanced Edge Detection
|
|
if config.edge_detection:
|
|
# Detect edges using Canny
|
|
edges = cv2.Canny(binary_mask, 100, 200)
|
|
|
|
# Create kernel based on edge_thickness
|
|
kernel = cv2.getStructuringElement(
|
|
cv2.MORPH_ELLIPSE,
|
|
(config.edge_thickness, config.edge_thickness)
|
|
)
|
|
|
|
# Dilate edges
|
|
edges = cv2.dilate(edges, kernel)
|
|
|
|
# Convert edge color from grayscale to color mask
|
|
edge_mask = np.zeros((*binary_mask.shape, 4), dtype=np.uint8)
|
|
edge_mask[edges > 0] = config.edge_color # Edge color from config
|
|
|
|
# Blend edges with original mask
|
|
binary_mask = cv2.addWeighted(
|
|
binary_mask,
|
|
0.7,
|
|
edges,
|
|
0.3,
|
|
0
|
|
)
|
|
|
|
# Handle mask expansion
|
|
if config.mask_expansion != 0:
|
|
kernel_size = abs(config.mask_expansion) // 10 + 1 # Scale factor for iteration count
|
|
kernel = np.ones((2, 2), np.uint8)
|
|
if config.mask_expansion > 0:
|
|
binary_mask = cv2.dilate(binary_mask, kernel, iterations=kernel_size)
|
|
else:
|
|
binary_mask = cv2.erode(binary_mask, kernel, iterations=kernel_size)
|
|
|
|
# Apply blur for smoother edges
|
|
if config.mask_blur > 0:
|
|
blur_amount = min(config.mask_blur, 20) # Limit maximum blur
|
|
blur_kernel_size = blur_amount * 2 + 1 # Ensure odd kernel size
|
|
binary_mask = cv2.GaussianBlur(
|
|
binary_mask,
|
|
(blur_kernel_size, blur_kernel_size),
|
|
0
|
|
)
|
|
|
|
# Edge refinement for better quality edges
|
|
if config.edge_refinement:
|
|
# Create erosion and dilation to identify edge areas
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
eroded = cv2.erode(binary_mask, kernel, iterations=1)
|
|
dilated = cv2.dilate(binary_mask, kernel, iterations=1)
|
|
|
|
# Edge mask (areas that differ between erosion and dilation)
|
|
edge_mask = dilated - eroded
|
|
|
|
# Apply stronger blur only to edge areas
|
|
strong_blur = cv2.GaussianBlur(binary_mask, (9, 9), 0)
|
|
|
|
# Blend the edge areas with strong blur, keep the rest from original mask
|
|
binary_mask = np.where(edge_mask > 0, strong_blur, binary_mask)
|
|
|
|
if config.remove_small_regions:
|
|
# Remove small disconnected areas
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary_mask, 8, cv2.CV_32S)
|
|
|
|
# Find the largest component (background is usually labeled as 0)
|
|
largest_label = 1 # Default to first component if no larger one is found
|
|
largest_size = 0
|
|
|
|
for i in range(1, num_labels):
|
|
if stats[i, cv2.CC_STAT_AREA] > largest_size:
|
|
largest_size = stats[i, cv2.CC_STAT_AREA]
|
|
largest_label = i
|
|
|
|
# Keep only regions larger than the configured size
|
|
clean_mask = np.zeros_like(binary_mask)
|
|
for i in range(1, num_labels):
|
|
if stats[i, cv2.CC_STAT_AREA] > config.small_region_size:
|
|
clean_mask[labels == i] = 255
|
|
binary_mask = clean_mask
|
|
|
|
if config.invert_generated_mask:
|
|
binary_mask = 255 - binary_mask
|
|
|
|
return Image.fromarray(binary_mask.astype(np.uint8), 'L')
|
|
|
|
except Exception as e:
|
|
logger.error(f"Mask refinement failed: {str(e)}")
|
|
return mask
|
|
|
|
@dataclass
|
|
class ChromaKeyConfig:
|
|
"""Configuration for chroma key processing"""
|
|
chroma_key_color: Union[str, Tuple[int, int, int]] = "green"
|
|
chroma_key_tolerance: float = 0.1
|
|
spill_reduction: float = 0.0
|
|
edge_refinement: bool = False
|
|
|
|
def remove_background_chroma(image: Image.Image, config: ChromaKeyConfig) -> Image.Image:
|
|
"""Enhanced chroma key with spill reduction and edge refinement"""
|
|
try:
|
|
# Convert to numpy array
|
|
img_np = np.array(image)
|
|
original_img = img_np.copy()
|
|
if img_np.shape[2] == 4:
|
|
img_np = img_np[:, :, :3]
|
|
|
|
# Convert to HSV color space
|
|
hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
|
|
|
|
# Basic color ranges
|
|
tolerance = int(30 * config.chroma_key_tolerance)
|
|
saturation_min = 30
|
|
value_min = 30
|
|
|
|
# Color-specific ranges
|
|
if config.chroma_key_color == "green":
|
|
lower = np.array([55 - tolerance, saturation_min, value_min])
|
|
upper = np.array([65 + tolerance, 255, 255])
|
|
elif config.chroma_key_color == "blue":
|
|
lower = np.array([110 - tolerance, saturation_min, value_min])
|
|
upper = np.array([130 + tolerance, 255, 255])
|
|
else: # red
|
|
lower1 = np.array([0, saturation_min, value_min])
|
|
upper1 = np.array([tolerance, 255, 255])
|
|
lower2 = np.array([180 - tolerance, saturation_min, value_min])
|
|
upper2 = np.array([180, 255, 255])
|
|
|
|
mask1 = cv2.inRange(hsv, lower1, upper1)
|
|
mask2 = cv2.inRange(hsv, lower2, upper2)
|
|
mask = cv2.bitwise_or(mask1, mask2)
|
|
|
|
# Create mask for non-red colors
|
|
if config.chroma_key_color != "red":
|
|
mask = cv2.inRange(hsv, lower, upper)
|
|
|
|
# Add spill suppression
|
|
if config.spill_reduction > 0:
|
|
# Create a wider range for spill detection
|
|
spill_tolerance = int(tolerance * 1.5)
|
|
|
|
if config.chroma_key_color == "green":
|
|
spill_lower = np.array([55 - spill_tolerance, saturation_min//2, value_min])
|
|
spill_upper = np.array([65 + spill_tolerance, 255, 255])
|
|
spill_mask = cv2.inRange(hsv, spill_lower, spill_upper)
|
|
elif config.chroma_key_color == "blue":
|
|
spill_lower = np.array([110 - spill_tolerance, saturation_min//2, value_min])
|
|
spill_upper = np.array([130 + spill_tolerance, 255, 255])
|
|
spill_mask = cv2.inRange(hsv, spill_lower, spill_upper)
|
|
else: # red
|
|
spill_lower1 = np.array([0, saturation_min//2, value_min])
|
|
spill_upper1 = np.array([spill_tolerance, 255, 255])
|
|
spill_lower2 = np.array([180 - spill_tolerance, saturation_min//2, value_min])
|
|
spill_upper2 = np.array([180, 255, 255])
|
|
spill_mask1 = cv2.inRange(hsv, spill_lower1, spill_upper1)
|
|
spill_mask2 = cv2.inRange(hsv, spill_lower2, spill_upper2)
|
|
spill_mask = cv2.bitwise_or(spill_mask1, spill_mask2)
|
|
|
|
# Find transition areas (spill areas not fully keyed out)
|
|
transition_mask = cv2.bitwise_and(spill_mask, cv2.bitwise_not(mask))
|
|
|
|
# Apply spill reduction to transition areas
|
|
if np.any(transition_mask > 0):
|
|
# Convert RGB to grayscale only for spill areas
|
|
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
|
|
|
# Create an image where spill areas are desaturated
|
|
desaturated_img = img_np.copy()
|
|
|
|
# Apply weighted desaturation based on spill_reduction setting
|
|
desaturation_strength = config.spill_reduction
|
|
transition_pixels = (transition_mask > 0)
|
|
|
|
# Replace color channels with grayscale weighted by strength
|
|
for c in range(3):
|
|
desaturated_img[:, :, c] = np.where(
|
|
transition_pixels,
|
|
img_np[:, :, c] * (1 - desaturation_strength) + gray * desaturation_strength,
|
|
img_np[:, :, c]
|
|
)
|
|
|
|
# Use the desaturated image for further processing
|
|
img_np = desaturated_img
|
|
|
|
# Invert mask (0 for background, 255 for foreground)
|
|
mask = 255 - mask
|
|
|
|
# Edge refinement
|
|
if config.edge_refinement:
|
|
# Create eroded and dilated versions to identify edge regions
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
mask_eroded = cv2.erode(mask, kernel, iterations=1)
|
|
mask_dilated = cv2.dilate(mask, kernel, iterations=1)
|
|
|
|
# Edge regions are the difference between dilated and eroded masks
|
|
edge_regions = mask_dilated - mask_eroded
|
|
|
|
# Apply strong blur only to edge regions
|
|
blurred_mask = cv2.GaussianBlur(mask, (9, 9), 0)
|
|
|
|
# Combine: use blurred mask for edge regions, original mask elsewhere
|
|
refined_mask = np.where(edge_regions > 0, blurred_mask, mask)
|
|
mask = refined_mask
|
|
|
|
# Return the processed mask
|
|
return Image.fromarray(mask, 'L')
|
|
|
|
except Exception as e:
|
|
logger.error(f"Chroma key removal failed: {str(e)}")
|
|
raise RuntimeError(f"Chroma key removal failed: {str(e)}") |