Spaces:
Sleeping
Sleeping
File size: 5,796 Bytes
35ef092 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
# PortraitBlurrer.py
import cv2
import numpy as np
from PIL import Image
class PortraitBlurrer:
def __init__(self, max_blur=31, depth_threshold=120,
feather_strength=3, sharpen_strength=1):
self.max_blur = max_blur
# Ensure max_blur is odd and positive
if self.max_blur % 2 == 0:
self.max_blur += 1
if self.max_blur <= 0:
self.max_blur = 3 # Default odd positive
self.depth_threshold = depth_threshold
self.feather_strength = feather_strength
self.sharpen_strength = sharpen_strength
def refine_depth_map(self, depth_map):
# Apply a bilateral filter to smooth depth while preserving edges
refined_depth = cv2.bilateralFilter(depth_map, 9, 75, 75)
return refined_depth
def create_subject_mask(self, depth_map):
_, mask = cv2.threshold(depth_map, self.depth_threshold, 255, cv2.THRESH_BINARY)
kernel = np.ones((5, 5), np.uint8)
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
ksize = self.feather_strength
if ksize % 2 == 0:
ksize += 1
if ksize <= 0:
ksize = 3
mask = cv2.GaussianBlur(mask, (ksize, ksize), 0)
return mask.astype(np.float32) / 255.0
def sharpen_image(self, image):
# Ensure sharpen_strength is not zero to avoid division issues later if needed
strength = max(0.1, self.sharpen_strength) # Prevent zero strength
# Simple sharpening kernel
kernel = np.array([[-1, -1, -1],
[-1, 9, -1],
[-1, -1, -1]])
# Apply the kernel - adjust strength application if needed
# A common way is to blend the sharpened with original based on strength
sharpened = cv2.filter2D(image, -1, kernel)
# Blend sharpened and original based on strength
# strength=1 means mostly sharpened, strength close to 0 means mostly original
if strength != 1.0: # Avoid unnecessary work if strength is 1
blended = cv2.addWeighted(image, 1.0 - (strength - 1.0) if strength > 1.0 else 1.0 ,
sharpened, strength if strength <= 1.0 else 1.0, 0)
# Basic clipping if values go out of range due to sharpening
return np.clip(blended, 0, 255).astype(np.uint8)
else:
# Basic clipping if values go out of range due to sharpening
return np.clip(sharpened, 0, 255).astype(np.uint8)
def apply_blur(self, original_bgr, depth_map_array):
# Resize depth map to match image dimensions
depth_resized = cv2.resize(depth_map_array, (original_bgr.shape[1], original_bgr.shape[0]),
interpolation=cv2.INTER_LINEAR)
refined_depth = self.refine_depth_map(depth_resized)
mask = self.create_subject_mask(refined_depth) # Float mask [0, 1]
blurred = cv2.GaussianBlur(original_bgr, (self.max_blur, self.max_blur), 0)
# Only sharpen if strength is significant
if self.sharpen_strength > 0.05: # Threshold to avoid unnecessary computation
sharpened_original = self.sharpen_image(original_bgr)
# Blend sharpened subject with original based on mask
foreground = sharpened_original * mask[:, :, np.newaxis] + \
original_bgr * (1 - mask[:, :, np.newaxis])
else:
foreground = original_bgr # Use original if no sharpening
# Blend the (potentially sharpened) foreground with the blurred background
background = blurred * (1 - mask[:, :, np.newaxis])
# Combine the foreground (where mask is 1) and background (where mask is 0)
# Note: Foreground already contains the original where it wasn't sharpened
# A potentially better blend:
result = original_bgr * mask[:, :, np.newaxis] + blurred * (1 - mask[:, :, np.newaxis])
if self.sharpen_strength > 0.05:
sharpened_subject_only = self.sharpen_image(original_bgr)
# Apply sharpening only where the mask is high
result = sharpened_subject_only * mask[:, :, np.newaxis] + result * (1 - mask[:, :, np.newaxis])
# Ensure result is uint8
final_result = np.clip(result, 0, 255).astype(np.uint8)
# Return the final blurred image as a NumPy array (BGR)
# Also return the refined depth map and the mask for potential display
return final_result, refined_depth, (mask * 255).astype(np.uint8)
def process_image(self, original_bgr_np, depth_image_pil):
depth_map_array = np.array(depth_image_pil)
if len(depth_map_array.shape) > 2:
# Assuming input PIL depth map might be RGB, convert to grayscale
depth_map_array = cv2.cvtColor(depth_map_array, cv2.COLOR_RGB2GRAY)
elif len(depth_map_array.shape) == 2:
# Already grayscale, ensure it's uint8 if necessary (though pipeline likely outputs it correctly)
if depth_map_array.dtype != np.uint8:
# Normalize if it's float or other types before potential processing
if depth_map_array.max() > 1.0: # Basic check if it might be 0-255
depth_map_array = depth_map_array.astype(np.uint8)
else: # Assume 0-1 float, scale to 0-255
depth_map_array = (depth_map_array * 255).astype(np.uint8)
# apply_blur now returns the result, depth map, and mask
blurred_image_np, refined_depth_np, mask_np = self.apply_blur(original_bgr_np, depth_map_array)
# Return the blurred image, the refined depth map (grayscale), and the mask (grayscale)
return blurred_image_np, refined_depth_np, mask_np |