agrobot_base/Python/heatmap_detect/v1.py

70 lines
2.4 KiB
Python
Raw Normal View History

2024-08-13 16:28:41 +00:00
import cv2
import numpy as np
import matplotlib.pyplot as plt
# Carregar a imagem
image_path = "mapa_calor.jpg"
image = cv2.imread(image_path)
# Converter a imagem de BGR para RGB
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Definir intervalos de cores em HSV para segmentação
lower_blue = np.array([100, 150, 0])
upper_blue = np.array([140, 255, 255])
lower_green = np.array([40, 52, 72])
upper_green = np.array([80, 255, 255])
lower_red1 = np.array([0, 70, 50])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 70, 50])
upper_red2 = np.array([180, 255, 255])
# Converter a imagem para HSV
image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# Criar máscaras para cada cor
mask_blue = cv2.inRange(image_hsv, lower_blue, upper_blue)
mask_green = cv2.inRange(image_hsv, lower_green, upper_green)
mask_red1 = cv2.inRange(image_hsv, lower_red1, upper_red1)
mask_red2 = cv2.inRange(image_hsv, lower_red2, upper_red2)
mask_red = cv2.bitwise_or(mask_red1, mask_red2)
# Função para calcular a distância com base na cor
def get_distance_from_color(color):
blue_dist = 500
red_dist = 4000
dist = blue_dist + (red_dist - blue_dist) * (color[0] / 255)
return dist
# Função para processar a máscara e encontrar o maior contorno
def get_largest_contour(mask):
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
largest_contour = max(contours, key=cv2.contourArea)
return largest_contour
return None
# Encontrar o maior contorno em cada máscara
contour_blue = get_largest_contour(mask_blue)
contour_green = get_largest_contour(mask_green)
contour_red = get_largest_contour(mask_red)
# Desenhar caixas delimitadoras e calcular distâncias
output_image = image_rgb.copy()
object_details = []
for contour, color in zip([contour_blue, contour_green, contour_red], [(0, 0, 255), (0, 255, 0), (255, 0, 0)]):
if contour is not None:
x, y, w, h = cv2.boundingRect(contour)
color_center = image_rgb[y + h // 2, x + w // 2]
distance = get_distance_from_color(color_center)
object_details.append((x, y, w, h, distance))
cv2.rectangle(output_image, (x, y), (x + w, y + h), (255, 255, 255), 2)
cv2.putText(output_image, f'D: {distance:.0f}mm', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
# Mostrar a imagem processada
plt.figure(figsize=(10, 10))
plt.imshow(output_image)
plt.axis('off')
plt.show()