104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
from skimage.morphology import skeletonize
|
|
|
|
|
|
def show(title, img, scale=0.6):
|
|
resized = cv2.resize(img, (0, 0), fx=scale, fy=scale)
|
|
cv2.imshow(title, resized)
|
|
cv2.waitKey(0)
|
|
|
|
|
|
def extract_wire_mask(hsv_img, lower, upper, name='color'):
|
|
mask = cv2.inRange(hsv_img, lower, upper)
|
|
# show(f"{name} - begin", mask)
|
|
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
|
|
closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
|
dilated = cv2.dilate(closed, kernel, iterations=1)
|
|
# show(f"{name} - process + exbound", dilated)
|
|
return dilated
|
|
|
|
|
|
def get_skeleton(binary_mask, name='skeleton'):
|
|
skel = skeletonize(binary_mask // 255).astype(np.uint8) * 255
|
|
# show(f"{name} - bouns", skel)
|
|
return skel
|
|
|
|
|
|
def find_endpoints(skel_img):
|
|
endpoints = []
|
|
h, w = skel_img.shape
|
|
for y in range(1, h - 1):
|
|
for x in range(1, w - 1):
|
|
if skel_img[y, x] == 255:
|
|
patch = skel_img[y - 1:y + 2, x - 1:x + 2]
|
|
if cv2.countNonZero(patch) == 2:
|
|
endpoints.append((x, y))
|
|
return endpoints
|
|
|
|
|
|
def detect_wires_and_endpoints(image):
|
|
original = image
|
|
img = cv2.resize(original, (1000, int(original.shape[0] * 1000 / original.shape[1])))
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
|
|
ranges = {
|
|
'red': [(np.array([0, 112, 38]), np.array([8, 255, 255])),
|
|
(np.array([160, 70, 50]), np.array([180, 255, 255]))],
|
|
'green': [(np.array([35, 80, 80]), np.array([85, 255, 255]))],
|
|
'yellow': [(np.array([19, 115, 103]), np.array([35, 255, 255]))]
|
|
}
|
|
|
|
result_img = img.copy()
|
|
all_wires = []
|
|
|
|
for color_name, hsv_ranges in ranges.items():
|
|
# print(f"\n🟢 正在处理颜色: {color_name.upper()}")
|
|
|
|
mask_total = np.zeros(hsv.shape[:2], dtype=np.uint8)
|
|
for (lower, upper) in hsv_ranges:
|
|
mask = extract_wire_mask(hsv, lower, upper, color_name)
|
|
mask_total = cv2.bitwise_or(mask_total, mask)
|
|
|
|
skeleton = get_skeleton(mask_total, f"{color_name}_skeleton")
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(skeleton, connectivity=8)
|
|
|
|
for i in range(1, num_labels):
|
|
wire_mask = (labels == i).astype(np.uint8) * 255
|
|
pixel_count = cv2.countNonZero(wire_mask)
|
|
|
|
if pixel_count < 100:
|
|
continue
|
|
|
|
endpoints = find_endpoints(wire_mask)
|
|
wire_vis = cv2.cvtColor(wire_mask, cv2.COLOR_GRAY2BGR)
|
|
|
|
if len(endpoints) >= 2:
|
|
start, end = endpoints[0], endpoints[-1]
|
|
# print(f"✅ {color_name.upper()}导线 #{i}: 起点 {start}，终点 {end}，像素数 {pixel_count}")
|
|
cv2.circle(wire_vis, start, 6, (0, 0, 255), -1)
|
|
cv2.circle(wire_vis, end, 6, (255, 0, 0), -1)
|
|
cv2.line(wire_vis, start, end, (0, 255, 255), 2)
|
|
|
|
cv2.circle(result_img, start, 6, (0, 0, 255), -1)
|
|
cv2.circle(result_img, end, 6, (255, 0, 0), -1)
|
|
cv2.line(result_img, start, end, (0, 255, 255), 2)
|
|
|
|
# 保存导线数据
|
|
wire_data = {
|
|
"start": {"x": int(start[0]), "y": int(start[1])},
|
|
"end": {"x": int(end[0]), "y": int(end[1])},
|
|
}
|
|
all_wires.append(wire_data)
|
|
|
|
# show(f"{color_name.upper()} 导线 #{i}", wire_vis)
|
|
|
|
# 显示图像
|
|
# cv2.imshow('tmp', result_img)
|
|
# cv2.waitKey(0)
|
|
return all_wires
|
|
|
|
# 示例调用
|
|
# detect_wires_and_endpoints("img/5.jpg")
|