mirror of
https://github.com/Astatin3/fabric-point-cloud.git
synced 2026-06-08 16:18:07 -06:00
Astatin3
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import ktb
|
|
import mediapipe as mp
|
|
|
|
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!")
|
|
return create_person_mask(depth_map, results.pose_landmarks, rgb_image.shape)
|
|
else:
|
|
return np.zeros(rgb_image.shape[:2], dtype=np.uint8)
|
|
|
|
|
|
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)
|
|
|
|
# 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
|
|
|
|
|
|
# while running:
|
|
# # vis.update_geometry(reconstruction)
|
|
# # print("E")
|
|
# if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
# running = False
|
|
# running = vis.poll_events()
|
|
# vis.update_renderer()
|
|
|