Work on depth based point cloud masking

3dscan/skeleton.py

@@ -1,7 +1,10 @@
+import math
+
 import cv2
 import numpy as np
 import ktb
 import mediapipe as mp
+from collections import deque
 
 from scipy.spatial import ConvexHull
 
@@ -33,10 +36,117 @@ def calc_mask(depth_map, rgb_image):
     if results.pose_landmarks:
         # Create person mask
         print("Found person!")
+        # cv2.imshow("Test", generate_distance_map(results.pose_landmarks, rgb_image.shape))
         return create_person_mask(depth_map, results.pose_landmarks, rgb_image.shape)
     else:
         return np.zeros(rgb_image.shape[:2], dtype=np.uint8)
 
+def clamp(x, min_value, max_value):
+    return max(min(max_value, x), min_value)
+
+
+def get_blurred_point_value(image, point, kernel_size=(5, 5), sigma=1.0):
+    # Ensure odd kernel size
+    kernel_size = tuple(k + 1 if k % 2 == 0 else k for k in kernel_size)
+
+    # Calculate the half size of the kernel
+    half_width = kernel_size[0] // 2
+    half_height = kernel_size[1] // 2
+
+    # Extract the region of interest (ROI) around the point
+    x, y = point
+    roi = image[max(0, y - half_height):min(image.shape[0], y + half_height + 1),
+          max(0, x - half_width):min(image.shape[1], x + half_width + 1)]
+
+    # Apply Gaussian blur to the ROI
+    blurred_roi = cv2.GaussianBlur(roi, kernel_size, sigma)
+
+    # Get the value at the center of the blurred ROI
+    center_y, center_x = blurred_roi.shape[0] // 2, blurred_roi.shape[1] // 2
+    blurred_value = blurred_roi[center_y, center_x]
+
+    return blurred_value
+
+def gradient_line(mask, x1, y1, x2, y2, depth_1, depth_2, circle_dist=15, step_size=5):
+    circle_dist = 10
+
+    cv2.circle(mask, (x1, y1), circle_dist * 2, int(depth_1), thickness=-1)
+    # cv2.circle(mask, (x2, y2), circle_dist * 2, int(depth_2), thickness=-1)
+
+    # print(x1, y1, x2, y2)
+
+    distance = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
+    step_count = int(distance / step_size)
+    print(step_count)
+
+    if step_count <= 0:
+        return mask
+
+    distance_x = (x2 - x1) / step_count
+    distance_y = (y2 - y1) / step_count
+
+    if distance_x == 0 or distance_y == 0:
+        return mask
+
+    color_step_distance = (depth_2 - depth_1) / step_size
+
+    for i in range(int(step_count)):
+        x = x1 + i * distance_x
+        y = y1 + i * distance_y
+        color = depth_1 + color_step_distance * i
+
+        tmp_mask = mask.copy()
+        cv2.circle(tmp_mask, (int(x), int(y)), circle_dist * 2, color, thickness=-1)
+        mask = cv2.addWeighted(mask, 0.5, tmp_mask, 0.5, 1.0)
+        del tmp_mask
+    return mask
+
+SKELETON_DEPTH_TO_KINECT = 4096
+
+def generate_distance_map(depth_map, pose_landmarks, image_shape, circle_dist=15, step_size=10):
+    # Create a black background
+    # mask = np.zeros(image_shape[:2], dtype=np.uint8)
+    mask = depth_map
+
+    connections = []
+
+    # Draw pose landmarks and connections
+    for connection in mp_pose.POSE_CONNECTIONS:
+        start_point = pose_landmarks.landmark[connection[0]]
+        end_point = pose_landmarks.landmark[connection[1]]
+
+        x1, y1 = int(clamp(start_point.x, 0, .999) * image_shape[1]), int(clamp(start_point.y, 0, .999) * image_shape[0])
+        x2, y2 = int(clamp(end_point.x, 0, .999) * image_shape[1]), int(clamp(end_point.y, 0, .999) * image_shape[0])
+
+        # print(start_point.z, end_point.z)
+
+        # depth_1 = start_point.z * SKELETON_DEPTH_TO_KINECT
+        # depth_2 = end_point.z * SKELETON_DEPTH_TO_KINECT
+
+        depth_1 = get_blurred_point_value(depth_map, (x1, y1), kernel_size=(15, 15))
+        depth_2 = get_blurred_point_value(depth_map, (x2, y2), kernel_size=(15, 15))
+
+        connections.append({
+            "x1": x1, "y1": y1, "x2":x2, "y2":y2, "depth_1":depth_1, "depth_2":depth_2
+        })
+
+    n = len(connections)  # Get the length of the array
+    if n > 1:
+        for i in range(1, n):  # Iterate over the array starting from the second element
+            key = connections[i]  # Store the current element as the key to be inserted in the right position
+            j = i - 1
+            while j >= 0 and key["depth_1"] > connections[j]["depth_1"]:  # Move elements greater than key one position ahead
+                connections[j + 1] = connections[j]  # Shift elements to the right
+                j -= 1
+            connections[j + 1] = key  # Insert the key in the correct position
+
+
+    for c in connections:
+        mask = gradient_line(mask, c["x1"], c["y1"], c["x2"], c["y2"], c["depth_1"], c["depth_2"], circle_dist)
+
+
+    return mask
+
 
 def create_person_mask(depth_map, pose_landmarks, image_shape, distance_threshold=25):
     # Create an empty mask
@@ -55,25 +165,11 @@ def create_person_mask(depth_map, pose_landmarks, image_shape, distance_threshol
     # Dilate the mask to include nearby pixels
     kernel = np.ones((distance_threshold, distance_threshold), np.uint8)
     mask = cv2.dilate(mask, kernel, iterations=1)
+
+    mask_2 = generate_distance_map(depth_map, pose_landmarks, image_shape, distance_threshold)
+    cv2.imshow("mask_2", (mask_2 / 2048)*mask)
     
-    # Get depth values for the pose landmarks
-    landmark_depths = []
-    for landmark in pose_landmarks.landmark:
-        x, y = int(landmark.x * image_shape[1]), int(landmark.y * image_shape[0])
-        if 0 <= x < image_shape[1] and 0 <= y < image_shape[0]:
-            landmark_depths.append(depth_map[y, x])
-    
-    # Calculate depth range
-    min_depth = np.percentile(landmark_depths, 0)  # 5th percentile to avoid outliers
-    max_depth = np.percentile(landmark_depths, 100)  # 95th percentile to avoid outliers
-    
-    # Refine the mask using depth information
-    depth_mask = (depth_map >= min_depth) & (depth_map <= max_depth)
-    
-    # Combine the initial mask with the depth mask
-    final_mask = mask & depth_mask
-    
-    return final_mask
+    return mask
 
 
 # while running: