Create depth_segmentation.py

depth_segmentation.py

@@ -0,0 +1,43 @@
+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