fabric-point-cloud/3dscan/skeleton.py

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()