Files
AutoControlSystem-git/Vision/tool/utils.py
2024-09-18 18:31:08 +08:00

241 lines
7.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project AutoControlSystem-master
@File utils.py
@IDE PyCharm
@Author hjw
@Date 2024/8/29 15:07
'''
import numpy as np
import cv2
import psutil
from psutil._common import bytes2human
from sklearn.decomposition import PCA
def pca(points):
# 计算PCA
mean = np.mean(points, axis=0)
cov = np.cov(points, rowvar=False)
eigenvalues, eigenvectors = np.linalg.eigh(cov)
# 对特征值进行排序并找到对应的特征向量
sorted_indices = np.argsort(eigenvalues)[::-1]
sorted_eigenvalues = eigenvalues[sorted_indices]
sorted_eigenvectors = eigenvectors[:, sorted_indices]
# # 打印PCA结果
# print("PCA Mean:", mean)
# print("PCA Eigenvalues:", sorted_eigenvalues)
# print("PCA Eigenvectors:\n", sorted_eigenvectors)
# 通常,最小特征值对应的特征向量是平面的法线
normal_vector = sorted_eigenvectors[:, 0] if sorted_eigenvalues[0] < sorted_eigenvalues[
1] else sorted_eigenvectors[:, 1]
# print("法向量:", normal_vector)
return normal_vector
def compute_pca_direction(points):
# points 是物体点云的坐标
pca = PCA(n_components=3)
pca.fit(points)
primary_direction = pca.components_[0] # 第一主成分对应的方向向量
return primary_direction
def remove_nan(pm, y, x):
piont_x, piont_y, piont_z = pm[y, x]
if np.isnan(piont_x):
for i in range(10):
piont_x, piont_y, piont_z = pm[y+i, x]
if np.isnan(piont_x)==False:
break
return piont_x, piont_y, piont_z
def get_disk_space(path='C:'):
usage = psutil.disk_usage(path)
space_free = bytes2human(usage.free)
# space_total = bytes2human(usage.total)
# space_used = bytes2human(usage.used)
# space_free = bytes2human(usage.free)
# space_used_percent = bytes2human(usage.percent)
space_free = float(space_free[:-1])
return space_free
def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=True):
if first_depth == True:
sorted_id = sorted(range(len(Depth_Z)), key=lambda k: Depth_Z[k], reverse=False)
Depth_Z1 = [Depth_Z[i] for i in sorted_id]
RegionalArea1 = [RegionalArea[i] for i in sorted_id]
for i in range(len(Depth_Z1)):
if RegionalArea1[i] > RegionalArea_Threshold:
return sorted_id[i]
else:
sorted_id = sorted(range(len(RegionalArea)), key=lambda k: RegionalArea[k], reverse=True)
Depth_Z1 = [Depth_Z[i] for i in sorted_id]
RegionalArea1 = [RegionalArea[i] for i in sorted_id]
for i in range(len(Depth_Z1)):
if RegionalArea1[i] > RegionalArea_Threshold:
return sorted_id[i]
class_names = ['box', 'other']
# Create a list of colors for each class where each color is a tuple of 3 integer values
rng = np.random.default_rng(3)
colors = rng.uniform(0, 255, size=(len(class_names), 3))
def nms(boxes, scores, iou_threshold):
# Sort by score
sorted_indices = np.argsort(scores)[::-1]
keep_boxes = []
while sorted_indices.size > 0:
# Pick the last box
box_id = sorted_indices[0]
keep_boxes.append(box_id)
# Compute IoU of the picked box with the rest
ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
# Remove boxes with IoU over the threshold
keep_indices = np.where(ious < iou_threshold)[0]
# print(keep_indices.shape, sorted_indices.shape)
sorted_indices = sorted_indices[keep_indices + 1]
return keep_boxes
def compute_iou(box, boxes):
# Compute xmin, ymin, xmax, ymax for both boxes
xmin = np.maximum(box[0], boxes[:, 0])
ymin = np.maximum(box[1], boxes[:, 1])
xmax = np.minimum(box[2], boxes[:, 2])
ymax = np.minimum(box[3], boxes[:, 3])
# Compute intersection area
intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
# Compute union area
box_area = (box[2] - box[0]) * (box[3] - box[1])
boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
union_area = box_area + boxes_area - intersection_area
# Compute IoU
iou = intersection_area / union_area
return iou
def xywh2xyxy(x):
# Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
y = np.copy(x)
y[..., 0] = x[..., 0] - x[..., 2] / 2
y[..., 1] = x[..., 1] - x[..., 3] / 2
y[..., 2] = x[..., 0] + x[..., 2] / 2
y[..., 3] = x[..., 1] + x[..., 3] / 2
return y
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def draw_detections(image, boxes, scores, class_ids, mask_alpha=0.3, mask_maps=None):
img_height, img_width = image.shape[:2]
size = min([img_height, img_width]) * 0.0006
text_thickness = int(min([img_height, img_width]) * 0.001)
mask_img = draw_masks(image, boxes, class_ids, mask_alpha, mask_maps)
# Draw bounding boxes and labels of detections
for box, score, class_id in zip(boxes, scores, class_ids):
color = colors[class_id]
x1, y1, x2, y2 = box.astype(int)
# Draw rectangle
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 2)
label = class_names[class_id]
caption = f'{label} {int(score * 100)}%'
(tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=size, thickness=text_thickness)
th = int(th * 1.2)
cv2.rectangle(mask_img, (x1, y1),
(x1 + tw, y1 - th), color, -1)
cv2.putText(mask_img, caption, (x1, y1),
cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
return mask_img
def draw_masks(image, boxes, class_ids, mask_alpha=0.3, mask_maps=None):
mask_img = image.copy()
# Draw bounding boxes and labels of detections
for i, (box, class_id) in enumerate(zip(boxes, class_ids)):
color = colors[class_id]
x1, y1, x2, y2 = box.astype(int)
# Draw fill mask image
if mask_maps is None:
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
else:
crop_mask = mask_maps[i][y1:y2, x1:x2, np.newaxis]
crop_mask_img = mask_img[y1:y2, x1:x2]
crop_mask_img = crop_mask_img * (1 - crop_mask) + crop_mask * color
mask_img[y1:y2, x1:x2] = crop_mask_img
return cv2.addWeighted(mask_img, mask_alpha, image, 1 - mask_alpha, 0)
def draw_comparison(img1, img2, name1, name2, fontsize=2.6, text_thickness=3):
(tw, th), _ = cv2.getTextSize(text=name1, fontFace=cv2.FONT_HERSHEY_DUPLEX,
fontScale=fontsize, thickness=text_thickness)
x1 = img1.shape[1] // 3
y1 = th
offset = th // 5
cv2.rectangle(img1, (x1 - offset * 2, y1 + offset),
(x1 + tw + offset * 2, y1 - th - offset), (0, 115, 255), -1)
cv2.putText(img1, name1,
(x1, y1),
cv2.FONT_HERSHEY_DUPLEX, fontsize,
(255, 255, 255), text_thickness)
(tw, th), _ = cv2.getTextSize(text=name2, fontFace=cv2.FONT_HERSHEY_DUPLEX,
fontScale=fontsize, thickness=text_thickness)
x1 = img2.shape[1] // 3
y1 = th
offset = th // 5
cv2.rectangle(img2, (x1 - offset * 2, y1 + offset),
(x1 + tw + offset * 2, y1 - th - offset), (94, 23, 235), -1)
cv2.putText(img2, name2,
(x1, y1),
cv2.FONT_HERSHEY_DUPLEX, fontsize,
(255, 255, 255), text_thickness)
combined_img = cv2.hconcat([img1, img2])
if combined_img.shape[1] > 3840:
combined_img = cv2.resize(combined_img, (3840, 2160))
return combined_img