File size: 1,716 Bytes
180d473 |
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 |
import cv2
import numpy as np
from PIL import Image
def segment_image_by_depth(original_img_path, depth_map_path, num_segments=3):
# Load original image (JPG) and depth map
original_img = Image.open(original_img_path).convert("RGB") # Convert to RGB
depth_map = cv2.imread(depth_map_path, cv2.IMREAD_GRAYSCALE)
# Normalize depth map
depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
# Calculate depth ranges based on the number of segments
segment_ranges = np.linspace(0, 255, num_segments + 1).astype(int)
segmented_images = []
for i in range(num_segments):
# Generate mask for the current depth range
lower, upper = int(segment_ranges[i]), int(segment_ranges[i + 1])
mask = cv2.inRange(depth_map, lower, upper)
# Convert mask to 3 channels to match the original image
mask_rgb = cv2.merge([mask, mask, mask])
# Apply mask to the original image
segmented_img = cv2.bitwise_and(np.array(original_img), mask_rgb)
# Convert back to PIL Image and add an alpha channel
segmented_img_pil = Image.fromarray(segmented_img).convert("RGBA")
alpha_channel = Image.fromarray(mask).convert("L") # Use the mask as the alpha channel
segmented_img_pil.putalpha(alpha_channel) # Add transparency based on depth mask
segmented_images.append(segmented_img_pil)
# Save each segmented image as PNG
output_paths = []
for idx, seg_img in enumerate(segmented_images):
output_path = f"outputs/segment_{idx + 1}.png"
seg_img.save(output_path)
output_paths.append(output_path)
return output_paths
|