agrobot_base/Python/kinect-processor.py

102 lines
3.6 KiB
Python

import paho.mqtt.client as mqtt
import numpy as np
import cv2
import json
def on_connect(client, userdata, flags, rc):
print("Connected with result code " + str(rc))
client.subscribe("kinect/depth")
def on_message(client, userdata, msg):
depth_pixels = json.loads(msg.payload)
depth_values = [pixel['Depth'] for pixel in depth_pixels] # Extraia os valores de profundidade
heatmap, obstacles = process_depth_data(depth_values)
send_obstacles(client, obstacles)
def process_depth_data(depth_values):
# Converta os valores de profundidade para um array numpy
depth_array = np.array(depth_values, dtype=np.uint16)
# Redimensione o array para a forma correta, se necessário
depth_array = depth_array.reshape((480, 640)) # Exemplo de resolução 480x640, ajuste conforme necessário
# Gere o mapa de calor (heatmap)
heatmap = cv2.applyColorMap(cv2.convertScaleAbs(depth_array, alpha=0.03), cv2.COLORMAP_JET)
# Segmente a imagem para encontrar obstáculos
obstacles = segment_image_by_color(heatmap, depth_array)
return heatmap, obstacles
def segment_image_by_color(heatmap, depth_array):
# Defina a cor alvo e o limiar
target_color = np.array([0, 0, 255]) # Exemplo: vermelho puro
threshold = 50 # Ajuste conforme necessário
# Converta o mapa de calor para o espaço de cores HSV
hsv_heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2HSV)
target_hsv = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2HSV)[0][0]
# Defina os limites inferior e superior para a segmentação
lower_bound = np.array([target_hsv[0] - threshold, 100, 100])
upper_bound = np.array([target_hsv[0] + threshold, 255, 255])
# Crie uma máscara para a cor alvo
mask = cv2.inRange(hsv_heatmap, lower_bound, upper_bound)
# Encontre contornos na máscara
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
obstacles = []
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
obstacle = calculate_obstacle_properties(x, y, w, h, depth_array)
obstacles.append(obstacle)
return obstacles
def calculate_obstacle_properties(x, y, width, height, depth_array):
# Calcular a distância média do obstáculo
obstacle_region = depth_array[y:y+height, x:x+width]
distance = np.mean(obstacle_region)
# Calcular a relação pixels para milímetros
pixel_to_mm = calculate_pixel_to_mm_ratio(distance)
# Converter dimensões de pixels para milímetros
width_mm = width * pixel_to_mm
height_mm = height * pixel_to_mm
obstacle = {
"X": x,
"Y": y,
"Width": width,
"Height": height,
"Distance_mm": distance,
"Width_mm": width_mm,
"Height_mm": height_mm
}
return obstacle
def calculate_pixel_to_mm_ratio(distance):
# Exemplo de cálculo simplificado da relação pixels para milímetros
# A relação exata pode variar e pode precisar de ajustes
# Aqui assumimos uma relação linear simplificada
focal_length = 580 # Focal length do Kinect em pixels
sensor_width_mm = 70 # Largura do sensor em milímetros
pixel_to_mm = (distance / focal_length) * (sensor_width_mm / 640) * 10
return pixel_to_mm
def send_obstacles(client, obstacles):
message = json.dumps(obstacles)
client.publish("kinect/obstacles", message)
client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message
client.connect("localhost", 1883, 60) # Substitua pelo seu broker
client.loop_forever()