import os
import cv2
import json
import glob
import torch
import traceback
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import leastsq
def cal_border(image, model_size, stride):
shape = image.shape[:2] # current shape [height, width]
new_shape = (model_size, model_size)
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
unpad_size = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - unpad_size[0], new_shape[0] - unpad_size[1] # wh padding
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
dw /= 2 # divide padding into 2 sides
dh /= 2
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
return top, bottom, left, right, unpad_size
def split_face_tool_3d(depth_images, binarys, camera_config, rotor_config, invaild_value):
ret_dict = {
"0": 0,
"1": "Success",
"2": [], # Mask
"3": [], # DepthImage
"4": [], # ROI
"5": [], # Result Str
"6": [], # Result ID
}
try:
camera_config = json.loads(camera_config)
rotor_config = json.loads(rotor_config)
x_res = camera_config.get("3D相机X精度")
y_res = camera_config.get("3D相机Y精度")
z_res = camera_config.get("3D相机Z比例")
piece_h = rotor_config.get("单磁极侧壁弧长")
piece_w = rotor_config.get("侧壁宽度")
piece_num = rotor_config.get("磁极数")
peice_y_num = piece_h / y_res
ignore_y_pos = 5
if len(depth_images) != len(binarys):
raise Exception("深度图和二值化图输入数量不一致.")
for i in range(len(depth_images)):
depth_image = depth_images[i]
if depth_image.ndim == 3 and depth_image.shape[2] == 1:
depth_image = depth_image[:, :, 0]
binary = binarys[i]
result_str = []
result_id = []
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
min_area = piece_h * piece_w * 0.8 / (x_res * y_res)
valid_contours = [
c for c in contours
if cv2.contourArea(c) >= min_area and cv2.boundingRect(c)[1] > ignore_y_pos
]
valid_contours.sort(key=lambda c: cv2.boundingRect(c)[1])
if len(valid_contours) > piece_num:
valid_contours = valid_contours[:piece_num]
crop_contour_list = []
crop_contour_elevation_list = []
rect_list = []
for contour in valid_contours:
x, y, w, h = cv2.boundingRect(contour)
y_start = max(0, y - 5)
y_end = min(depth_image.shape[0], y + h + 10)
x_start = max(0, x)
x_end = min(depth_image.shape[1], x + w)
crop_h = y_end - y_start
crop_w = x_end - x_start
crop_contour = np.zeros((crop_h, crop_w), dtype=np.uint8)
offset_contour = contour - (x_start, y_start)
cv2.drawContours(crop_contour, [offset_contour], -1, 1, -1)
crop_elevation = depth_image[y_start:y_end, x_start:x_end]
crop_contour_elevation = np.where(crop_contour > 0, crop_elevation, invaild_value)
crop_contour_list.append(crop_contour)
crop_contour_elevation_list.append(crop_contour_elevation)
rect_list.append((x_start, y_start, crop_w, crop_h))
if len(rect_list) < piece_num:
result_str.append("侧面检测异常")
result_id.append(19)
else:
result_str.append("良品")
result_id.append(1)
for index, (x, y, w, h) in enumerate(rect_list):
ret_dict["2"].append({"image": crop_contour_list[index] * 255})
ret_dict["3"].append({"image": crop_contour_elevation_list[index]})
ret_dict["4"].append([[2.0, x, y, w, h, 0.0]])
col_sum = np.sum(crop_contour_list[index], axis=0)
if np.median(col_sum) > peice_y_num * 1.3:
result_str.append("硅钢片连接处拉胶丝")
result_id.append(9)
ret_dict["5"].append(result_str)
ret_dict["6"].append(result_id)
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = str(e)
return ret_dict
def depth_image_show(elevations, z_min, z_max, invaild_value, mode, show, z_res, show_3d, temp_path, code_path, up_dis, low_dis):
ret_dict = {
"0": 0,
"1": "Success",
"2": [],
"3": ""
}
try:
if len(elevations) == 0:
return ret_dict
if show_3d:
if code_path == "null" or temp_path == "null":
raise Exception("未输入路径")
if not os.path.exists(code_path) or not os.path.exists(temp_path):
raise Exception("路径不存在")
tif_files = glob.glob(os.path.join(temp_path, "*.npy"))
for f in tif_files:
os.remove(f)
cfg_dict = {
"z_min": z_min,
"z_max": z_max,
"invaild_value": invaild_value,
"auto_z": mode == "自适应计算高度范围"
}
with open(os.path.join(temp_path, "show_config.json"), "w") as f:
json.dump(cfg_dict, f)
with open(os.path.join(temp_path, "run.bat"), "w") as f:
f.write(r"conda activate py36 && python %s %s && exit"%(code_path, temp_path))
if show:
for index, elevation in enumerate(elevations):
if elevation.ndim == 3:
elevation = elevation[:, :, 0]
valid_mask = elevation > invaild_value
no_zero_data = elevation[valid_mask]
if mode == "自适应计算高度范围":
z_min = no_zero_data.min()*z_res - 10*z_res
z_max = no_zero_data.max()*z_res + 10*z_res
ret_dict["3"] += "z_min: %.2f z_max: %.2f z_mean: %.2f\n"%(z_min, z_max, np.mean(no_zero_data))
elif mode == "均值上下扩展":
mean = np.mean(no_zero_data)*z_res
z_min = mean - low_dis
z_max = mean + up_dis
ret_dict["3"] += "z_min: %.2f z_max: %.2f z_mean: %.2f\n"%(z_min, z_max, mean)
z_min = int(z_min/z_res)
z_max = int(z_max/z_res)
clip_elevation = np.clip(elevation, z_min, z_max)
clip_elevation[0, 0] = z_max
depth_normalized = cv2.normalize(clip_elevation, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
depth_normalized = depth_normalized.astype(np.uint8)
color_image = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_JET)
ret_dict["2"].append({"image": color_image})
if show_3d:
file_name = os.path.join(temp_path, "temp_%d.npy"%index)
np.save(file_name, elevation)
if show_3d:
command = "start /B %s"%os.path.join(temp_path, "run.bat")
os.system(command)
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
def get_contour_edge_points(contour_images, x_skip, y_skip, x_num, y_num):
ret_dict = {
"0": 0,
"1": "Success",
"2": [],
}
try:
for contour_image in contour_images:
h, w = contour_image.shape[:2]
contour_image = contour_image/255
x_step = int((w-2*x_skip)/(x_num-1)) if x_num != 1 else 0
y_step = int((np.sum(contour_image, axis=0).mean()-2*y_skip)/(y_num-1)) if y_num != 1 else 0
edge_points = []
for i in range(x_num):
x = x_skip + i*x_step
for i_y in range(h):
if contour_image[i_y, x] > 0:
break
for j in range(y_num):
edge_points.append([4.0, x, i_y + y_skip + j*y_step])
ret_dict["2"].append(edge_points)
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
def get_point_in_area(contour_elevations, edge_points_list, area_size, select_fun, invalid_value, min_valid_value, max_valid_value):
ret_dict = {
"0": 0,
"1": "Success",
"2": [],
"3": []
}
try:
for index, contour_elevation in enumerate(contour_elevations):
select_points = []
h, w = contour_elevation.shape[:2]
edge_points = edge_points_list[index]
if select_fun == "最小值":
for point in edge_points:
_, x, y = point
y0 = int(max(0, y-area_size/2))
y1 = int(min(y+area_size/2, h))
x0 = int(max(0, x-area_size/2))
x1 = int(min(x+area_size/2, w))
crop = contour_elevation[y0:y1, x0:x1]
invalid_mask = (crop==invalid_value) + (crop<min_valid_value) + (crop>max_valid_value)
arr_for_select = invalid_mask*(crop.max()+1) + crop*(1-invalid_mask)
min_index = np.argmin(arr_for_select)
min_index_2d = np.unravel_index(min_index, arr_for_select.shape)
select_points.append([float(min_index_2d[1]+x0), float(min_index_2d[0]+y0), arr_for_select[min_index_2d]])
elif select_fun == "最大值":
for point in edge_points:
_, x, y = point
y0 = int(max(0, y-area_size/2))
y1 = int(min(y+area_size/2, h))
x0 = int(max(0, x-area_size/2))
x1 = int(min(x+area_size/2, w))
crop = contour_elevation[y0:y1, x0:x1]
invalid_mask = (crop==invalid_value) + (crop<min_valid_value) + (crop>max_valid_value)
arr_for_select = invalid_mask*(crop.min()-1) + crop*(1-invalid_mask)
max_index = np.argmax(arr_for_select)
max_index_2d = np.unravel_index(max_index, arr_for_select.shape)
select_points.append([float(max_index_2d[1]+x0), float(max_index_2d[0]+y0), arr_for_select[max_index_2d]])
ret_dict["2"].append(select_points)
ret_dict["3"].append([[4.0, select_point[0], select_point[1]]for select_point in select_points])
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
def fit_func(p, x, y):
""" 数据拟合函数 """
a, b, c = p
return a * x + b * y + c
def residuals(p, x, y, z):
""" 误差函数 """
return z - fit_func(p, x, y)
def fit_plane(points_list, debug, x_res, y_res, z_res):
ret_dict = {
"0": 0,
"1": "Success",
"2": [],
"3": ""
}
try:
for points in points_list:
points = np.array(points, dtype=np.float32)
X = points[:, 0] * x_res
Y = points[:, 1] * y_res
Z = points[:, 2] * z_res
A = np.vstack([X, Y, np.ones(len(X))]).T
a, b, c = np.linalg.lstsq(A, Z, rcond=None)[0]
ret_dict["2"].append([a, b, c])
if debug:
residuals = Z - (a*X + b*Y + c)
ret_dict["3"] += f"Loss: {np.abs(residuals).max():.2f}\n"
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
def cal_plane_dis_gpu(elevations, plane_argvs, x_res, y_res, z_res, z_min_high, z_min_low):
ret_dict = {
"0": 0,
"1": "Success",
"2": [],
"3": [],
"4": []
}
try:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
use_gpu = device.type == 'cuda'
for index, elevation in enumerate(elevations):
if elevation.ndim == 3:
elevation = elevation[:, :, 0]
a, b, c = plane_argvs[index]
h, w = elevation.shape[:2]
elevation_tensor = torch.tensor(elevation, dtype=torch.float32, device=device)
x_coords = torch.arange(w, dtype=torch.float32, device=device) * x_res
y_coords = torch.arange(h, dtype=torch.float32, device=device) * y_res
X, Y = torch.meshgrid(x_coords, y_coords, indexing='xy')
z_plane = a * X + b * Y + c
elevation_mm = elevation_tensor * z_res
dis = elevation_mm - z_plane
valid_mask = elevation_tensor > 0
dis = torch.where(valid_mask, dis, -10000)
mask_high = ((dis >= z_min_high) & valid_mask).to(torch.uint8) * 255
mask_low = ((dis <= -z_min_low) & valid_mask).to(torch.uint8) * 255
dis_np = dis.cpu().numpy()
mask_high_np = mask_high.cpu().numpy()
mask_low_np = mask_low.cpu().numpy()
ret_dict["2"].append(dis_np)
ret_dict["3"].append(mask_high_np)
ret_dict["4"].append(mask_low_np)
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = str(e)
return ret_dict
def depth_show_z_value_tools(depth_images, draw_images, is_show, z_res, font_size, x_adj, y_adj, points_list):
ret_dict = {
"0": 0,
"1": "Success",
"2": []
}
try:
if not is_show:
return ret_dict
for index, depth_image in enumerate(depth_images):
points = points_list[index] if len(points_list) > index else []
draw = draw_images[index].copy()
for point in points:
_, x, y = point
x = int(x)
y = int(y)
value = depth_image[y, x, 0] * z_res
cv2.putText(draw, "%.3f"%value, (x+x_adj, y+y_adj), cv2.FONT_HERSHEY_SIMPLEX, font_size, (255, 255, 255), 2)
ret_dict["2"].append({"image": draw})
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
def cal_component_distance(depth_images, masks, seg_masks, show_images, expand_w, expand_h, dilate_size, z_res, cal_mode, default_depth):
ret_dict = {
"0": 0,
"1": "Success",
"2": [], # ROIs
"3": [], # Areas
"4": [] # Height
}
try:
if len(depth_images) == 0:
raise Exception("输入图像数为0.")
if len(depth_images) != len(masks) or len(depth_images) != len(seg_masks):
raise Exception("输入深度图和掩码、分割图数量不一致.")
for index, depth_image in enumerate(depth_images):
mask = masks[index]
seg_mask = seg_masks[index]
H, W = depth_image.shape[:2]
rois = []
areas = []
heights = []
contours, _ = cv2.findContours(seg_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
contour_areas = [(cv2.contourArea(contour), contour) for contour in contours]
contour_areas_sorted = sorted(contour_areas, key=lambda x: x[0], reverse=True)
for contour in contour_areas_sorted:
x, y, w, h = cv2.boundingRect(contour[1])
max_depth = depth_image[y:y+h, x:x+w].max()*z_res
crop_depth_image = depth_image[max(0, y-expand_h):min(y+h+expand_h, H), max(0, x-expand_w):min(x+w+expand_w, W)]
crop_mask = 1-mask[max(0, y-expand_h):min(y+h+expand_h, H), max(0, x-expand_w):min(x+w+expand_w, W)]/255
crop_seg_mask = seg_mask[max(0, y-expand_h):min(y+h+expand_h, H), max(0, x-expand_w):min(x+w+expand_w, W)]
valid_depth_value = crop_depth_image[np.where((crop_seg_mask+crop_mask)==0)]*z_res
if len(valid_depth_value) < expand_w*expand_h:
crop_depth_image = depth_image[max(0, y-2*expand_h):min(y+h+2*expand_h, H), max(0, x-2*expand_w):min(x+w+2*expand_w, W)]
crop_mask = 1-mask[max(0, y-2*expand_h):min(y+h+2*expand_h, H), max(0, x-2*expand_w):min(x+w+2*expand_w, W)]/255
crop_seg_mask = seg_mask[max(0, y-2*expand_h):min(y+h+2*expand_h, H), max(0, x-2*expand_w):min(x+w+2*expand_w, W)]
valid_depth_value = crop_depth_image[np.where((crop_seg_mask+crop_mask)==0)]*z_res
if len(valid_depth_value) < expand_w*expand_h:
plane_depth = default_depth
else:
if cal_mode == "均值":
plane_depth = np.mean(valid_depth_value)
height = max_depth - plane_depth
rois.append([2.0, x, y, w, h, 0.0])
areas.append(contour[0])
heights.append(height)
ret_dict["2"].append(rois)
ret_dict["3"].append(areas)
ret_dict["4"].append(heights)
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = str(traceback.format_exc())
return ret_dict
def restore_yolo_mask(images, masks, model_size, stride):
ret_dict = {
"0": 0,
"1": "Success",
"2": []
}
try:
for index, image in enumerate(images):
mask = masks[index]
top, bottom, left, right, unpad_size = cal_border(image, model_size, stride) # unpad_size: W*H
restore_size = (unpad_size[0] + left + right, unpad_size[1] + top + bottom)
mask = cv2.resize(mask, restore_size, interpolation=cv2.INTER_NEAREST)
real_mask = mask[top:mask.shape[1]-bottom, left:mask.shape[1]-right]
real_mask = cv2.resize(real_mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
ret_dict["2"].append({"image": real_mask})
except Exception as e:
ret_dict["0"] = 1
ret_dict["1"] = traceback.format_exc()
return ret_dict
优化以上代码,导入问题,第一次运行会出现卡死的情况,优化代码输出出来
最新发布