mirror of
https://github.com/Astatin3/fabric-point-cloud.git
synced 2026-06-09 08:38:04 -06:00
Astatin3
183 lines
6.1 KiB
Python
183 lines
6.1 KiB
Python
import math
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import ktb
|
|
import mediapipe as mp
|
|
from collections import deque
|
|
|
|
from scipy.spatial import ConvexHull
|
|
|
|
running = True
|
|
|
|
# Initialize MediaPipe Pose
|
|
mp_pose = mp.solutions.pose
|
|
mp_drawing = mp.solutions.drawing_utils
|
|
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)
|
|
|
|
|
|
def get_color():
|
|
return k.get_frame(ktb.COLOR)
|
|
|
|
def get_depth_map():
|
|
return k.get_frame(ktb.DEPTH)
|
|
|
|
def normalize_depth_map(depth_map):
|
|
min_val, max_val, _, _ = cv2.minMaxLoc(depth_map)
|
|
normalized = cv2.convertScaleAbs(depth_map, alpha=255.0/(max_val-min_val), beta=-min_val * 255.0/(max_val-min_val))
|
|
return normalized
|
|
|
|
def calc_mask(depth_map, rgb_image):
|
|
|
|
|
|
# Process the image and get pose landmarks
|
|
results = pose.process(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
|
|
mask = np.zeros(image_shape[:2], dtype=np.uint8)
|
|
|
|
# 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(start_point.x * image_shape[1]), int(start_point.y * image_shape[0])
|
|
x2, y2 = int(end_point.x * image_shape[1]), int(end_point.y * image_shape[0])
|
|
|
|
cv2.line(mask, (x1, y1), (x2, y2), 1, thickness=distance_threshold*2)
|
|
|
|
# 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)
|
|
|
|
return mask
|
|
|
|
|
|
# while running:
|
|
# # vis.update_geometry(reconstruction)
|
|
# # print("E")
|
|
# if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
# running = False
|
|
# running = vis.poll_events()
|
|
# vis.update_renderer()
|
|
|