|
import cv2 |
|
import numpy as np |
|
from PIL import Image |
|
|
|
def segment_image_by_depth(original_img_path, depth_map_path, num_segments=3): |
|
|
|
original_img = Image.open(original_img_path).convert("RGB") |
|
depth_map = cv2.imread(depth_map_path, cv2.IMREAD_GRAYSCALE) |
|
|
|
|
|
depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) |
|
|
|
|
|
segment_ranges = np.linspace(0, 255, num_segments + 1).astype(int) |
|
|
|
segmented_images = [] |
|
|
|
for i in range(num_segments): |
|
|
|
lower, upper = int(segment_ranges[i]), int(segment_ranges[i + 1]) |
|
mask = cv2.inRange(depth_map, lower, upper) |
|
|
|
|
|
mask_rgb = cv2.merge([mask, mask, mask]) |
|
|
|
|
|
segmented_img = cv2.bitwise_and(np.array(original_img), mask_rgb) |
|
|
|
|
|
segmented_img_pil = Image.fromarray(segmented_img).convert("RGBA") |
|
alpha_channel = Image.fromarray(mask).convert("L") |
|
segmented_img_pil.putalpha(alpha_channel) |
|
|
|
segmented_images.append(segmented_img_pil) |
|
|
|
|
|
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 |
|
|