import copy
import os
import cv2
# from Cython.Compiler.Naming import obj_base_cname
from shapely.creation import polygons
from ultralytics import YOLO
import numpy as np
from paddleocr import PaddleOCR
from rapidfuzz import process
import Levenshtein
import json
from collections import defaultdict
from shapely.geometry import Polygon, MultiPolygon, GeometryCollection,LineString
import pandas as pd
from shapely.errors import GEOSException
ocr = PaddleOCR(
use_angle_cls=True,
lang='en'
)
# indexc=0.6
def binarize_image(image):
"""
对输入的图像进行自适应阈值二值化处理
:param image: 输入的图像
:return: 二值化后的图像
"""
binary_image = cv2.adaptiveThreshold(image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY, 37, 3)
return binary_image
def sharpen_image(image, sharpen_strength=1.0):
"""
对输入的图像进行锐化处理，可控制锐化程度
:param image: 输入的图像
:param sharpen_strength: 锐化程度，默认值为 1.0
:return: 锐化后的图像
"""
kernel = np.array([[-1, -1, -1],
[-1, 9, -1],
[-1, -1, -1]])
sharpened = cv2.filter2D(image, -1, kernel * sharpen_strength)
return sharpened
def binarizeAndSharpenImage(image, sharpen_strength=1.0, do_sharpen=True):
"""
对指定路径的图像进行二值化和锐化处理
:param input_image_path: 输入图像的路径
:param output_image_path: 输出图像的路径
:param sharpen_strength: 锐化程度，默认值为 1.0
:param do_sharpen: 是否进行锐化处理，默认值为 True
"""
# 读取图片
# image = cv2.imread(input_image_path, cv2.IMREAD_GRAYSCALE)
# 二值化处理
binary_image = binarize_image(image)
if do_sharpen:
# 锐化处理
final_image = sharpen_image(binary_image, sharpen_strength)
else:
final_image = binary_image
return final_image
def image2OCR(im, conduits, features, path, index=0):
ol = []
for i, (point, center, angle, box, tag, points) in enumerate(conduits):
mask = np.zeros(im.shape[:2], dtype=np.uint8)
cv2.fillPoly(mask, [np.array(points, dtype=np.int32)], 255)
res = cv2.bitwise_and(im, im, mask=mask)
if angle > 20: angle += 270
res = rotateImage(res, angle)
sa = res.copy()
# 找到非黑色区域的边界
gray = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY) # 转为灰度图
_, binary = cv2.threshold(gray, 3, 255, cv2.THRESH_BINARY) # 将非黑色部分设为白色
# 计算非黑色区域的边界框
x, y, w, h = cv2.boundingRect(binary)
if w > 0 and h > 0:
# 裁剪非黑色部分
res = gray[y:y + h, x:x + w]
# cv2.imshow("ocr", res)
# cv2.waitKey(0)
sa = sa[y:y + h, x:x + w]
# 放大三倍
height, width = res.shape
if height > width:
res = rotateImage(res, 90)
sa = rotateImage(sa, 90)
# 增强对比度（直方图均衡化）
res = binarizeAndSharpenImage(res)
conduits[i].append(res)
if res is not None:
if tag == "empty":
ocrImage, match = np.ones((50, 700, 3), dtype=im.dtype) * 255, "empty"
ol.append(match)
continue
else:
sa = cv2.cvtColor(res, cv2.COLOR_GRAY2BGR)
# cv2.imwrite(f'{i}.jpg', res)
ocrImage, match = im2OCR(res, features, ol)
devices = {}
for key in features:
if key == "set": continue
for match in ol:
if match in features[key].values():
devices[key] = devices.get(key, 0) + 1
if not devices:
return None, None, None
device = max(devices, key=devices.get)
deviceList = features[device]["set"]
# ol = [item for item in ol if item in deviceList]
used_indices = set() # 用于记录已使用的索引
matches = []
for match in ol:
if match in deviceList:
# 找到第一个未使用的索引
for idx, item in enumerate(deviceList):
if item == match and idx not in used_indices:
matches.append([match, idx])
used_indices.add(idx) # 标记该索引为已使用
break
else:
# 如果找不到未使用的索引，返回默认值 99
matches.append([match, 99])
else:
matches.append([match, 99])
counters = [(i, value[1]) for i, value in enumerate(matches) if (value[1] != 99 and value[0] != "empty")]
if len(counters) < 2: return None, None, None #有效计数器数量检查
revs = [counters[i][1] - counters[i - 1][1] for i in range(1, len(counters))]
rev = 0
for item in revs:
if item > 0: rev += 1
else: rev -= 1
rev = True if rev > 0 else False
# if rev == 0: return None, None, None
s = 1 if rev else -1
# 查找最大连续段的起始索引和长度
max_start = None
max_length = 0
current_start = None
current_length = 0
max_end = 0
prev = max_end
for i, (index, value) in enumerate(counters):
step = counters[i - 1][1] + (counters[i][1] - counters[i - 1][1])
v = value
if i == 0 or counters[i - 1][1] + (counters[i][1] - counters[i - 1][1]) == value: # 判断是否连续
if current_start is None:
current_start = index
current_length += 1
prev = max_end
max_end = index
if current_length > max_length:
max_length = current_length
max_start = current_start
else:
current_start = index
current_length = 1
# # 检查最后一段是否是最长的
# if current_length > max_length:
# max_length = current_length
# max_start = current_start
if max_length < 2: return None, None, None #连续段长度要求
# length = len(features[device]["set"]) / 2
min_required = max(3, len(features[device]["set"]) / 8) #最小匹配数量 1/2->1/8 ，最少3个匹配
counter, pos = matches[1][1], 1
# counter, pos = matches[prev][1], prev
# if len(matches) < length: return None, None, None
if len(matches) < min_required:
print(f"匹配数量不足: {len(matches)} < {min_required}")
return None, None, None
return matches, device, [rev, counter, pos]
def image2Cls(image, model): # 图像分类工具
results = model.predict(image, verbose=False, save=False)
cls = results[0].probs.top1
names = [value for key, value in model.names.items()]
return names[int(cls)]
def rotateImage(image, angle, border_value=(0, 0, 0)): # 旋转图像 输出旋转后的图像
"""
随机旋转图像。
参数：
image: np.ndarray，输入图像（灰度或彩色均可）。
max_angle: float，最大旋转角度（度）。实际旋转角度在 [-max_angle, +max_angle] 随机选取。
border_value: tuple/int，旋转后边界填充值。灰度图可传单个值，彩色图可传三元组。
返回：
rotated: np.ndarray，旋转后的图像。
angle: float，实际旋转角度（度）。
"""
h, w = image.shape[:2]
center = (w / 2, h / 2)
# 计算旋转矩阵
M = cv2.getRotationMatrix2D(center, angle, 1.0)
# 计算旋转后图像的新尺寸，避免裁剪
cos = abs(M[0, 0])
sin = abs(M[0, 1])
new_w = int(h * sin + w * cos)
new_h = int(h * cos + w * sin)
# 调整平移部分，使图像居中
M[0, 2] += new_w / 2 - center[0]
M[1, 2] += new_h / 2 - center[1]
# 应用仿射变换
rotated = cv2.warpAffine(
image, M, (new_w, new_h),
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=border_value
)
return rotated
# wrong_pic_count = int(indexc*10)
def imageFlatten(image, model): # 图像方向校正工具
im = cv2.imread(image)
name = image2Cls(im, model)
if "180" in name: im = rotateImage(im, 180)
if name == "090": im = rotateImage(im, 270)
if name == "270": im = rotateImage(im, 90)
return im, name
def inverseTransformCoordinates(coords, angle, image_shape):
"""
将变换后的坐标逆变换回原始图像坐标。
参数：
coords: list of tuples，变换后的坐标列表 [(x1, y1), (x2, y2), ...]。
angle: int，图像旋转角度（顺时针方向）。
image_shape: tuple，图像的形状 (height, width)。
返回：
原始图像中的坐标列表。
"""
h, w = image_shape[:2]
nc = []
for xy in coords:
if angle == 90:
nc.append((int(xy[1]), int(w - xy[0])))
elif angle == 180:
nc.append((int(w - xy[0]), int(h- xy[1])))
elif angle == 270:
nc.append((int(h - xy[1]), int(xy[0])))
else:
nc.append((int(xy[0]), int(xy[1])))
return nc
def filter_objects_by_line(objects, image_shape, threshold=300):
"""
根据图像方向（横版或竖版）选择聚类轴向，并保留数量最多的一行或一列。
:param objects: 对象列表，每个对象结构为 [point, center, ...]，conduits
:param image_shape: 图像的形状 (height, width)，用于判断方向
:param threshold: 判断属于同一行/列的距离阈值（像素）
:return: 过滤后的对象列表
"""
if not objects:
return []
# 提取所有对象的中心点坐标
centers = [obj[1] for obj in objects]
# 判断图像方向，决定聚类轴向
height, width = image_shape[:2]
use_x = height > width # 竖图 → 按列聚类（x）
coord_index = 0 if use_x else 1
# 根据选定轴向进行聚类
clusters = defaultdict(list)
for idx, (x, y) in enumerate(centers): # enumerate:获取中心点
val = x if use_x else y
found = False
for cluster_id, cluster_points in clusters.items():
comp_val = cluster_points[0][coord_index]
if abs(val - comp_val) <= threshold:
clusters[cluster_id].append((x, y))
found = True
break
if not found:
clusters[len(clusters)] = [(x, y)]
# 找出最大簇
largest_cluster = max(clusters.values(), key=len, default=[])
# 获取对应簇的索引
filtered_indices = [i for i, c in enumerate(centers) if c in largest_cluster]
return [objects[i] for i in filtered_indices]
def image2Align(image, orientedClsModel, SegModel, noseClsModel):# 对输入图像进行方向校正和语义分割
img, tag = imageFlatten(image, orientedClsModel) # 图像方向校正
textImg = img.copy()
results = SegModel.predict(img, verbose=False, save=True, show_conf=False, show_boxes=True,
exist_ok=True, conf=0.65) # 语义分割
# img = results[0].plot(labels=False, boxes=False)
names = [value for key, value in SegModel.names.items()]
# cv2.imshow("img", results[0].plot(labels=False, boxes=False))
yList = []
contrast = 0
for res in results[0]:
xy = [[item[0], item[1]] for mask in res.masks.xy for item in mask]
yList.append(xy)
ns = [names[int(ite)] for ite in results[0].boxes.cls.cpu().tolist()]
zipList = [[xy, n] for xy, n in zip(yList, ns)]
clsDict = {}
for xy, key in zipList:
if key in clsDict: clsDict[key].append(xy)
else: clsDict[key] = [xy]
return img, textImg, clsDict, results, int(tag)
# img: 校正后的图像 textImg: 用于OCR的文本图像 clsDict: 分类结果字典 results: 分割结果 rotate: 旋转角度
def praseExcel(excelFile): # 从Excel文件中提取特征信息
features = {
"set": []
}
# 读取 Excel 文件
if excelFile.endswith((".xlsx", "xls")):
# print("Extracting " + excelFile)
geoDF = pd.read_excel(excelFile, sheet_name=0, header=0) # 这里可以指定 sheet 名称，如果不指定，默认读取第一个 sheet
for index, row in geoDF.iterrows():
# serial = row.iloc[9] #新台账
# device = row.iloc[2]
# code = row.iloc[10]
serial = row.iloc[0]
device = row.iloc[1]
code = row.iloc[2]
if type(code) is float: code = "empty"
# if type(code) is float or code =="N/A" : code = "empty" #新台账
if device in features:
features[device].update({serial: code})
features[device]["set"].append(code)
else:
features.update({device: {serial: code,"set": [code,]}})
features["set"].append(code)
return features
def im2OCR(gray, features, ol): # 识别图像中的文本，并将进行匹配
try:
# 进行 OCR 识别
if len(gray.shape) == 2: # 如果是灰度图像
# 将灰度图转换为三通道 BGR 格式
img_bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
result = ocr.predict(img_bgr)
else:
result = ocr.predict(gray)
# print(f"OCR result: {result[0]['rec_texts']}")
# 检查结果是否为空
if not result or len(result) == 0 or len(result[0]) == 0 or not result[0]['rec_texts']:
match = "Unidentified"
else:
# 提取识别到的文本
text = result[0]['rec_texts'][0].strip()
# 替换字符
text = text.replace("D", "0")
# 获取特征集中的代码
codes = features["set"]
# 使用 fuzzywuzzy 进行模糊匹配
matchList = process.extract(text, codes, limit=16)
best_code, best_dist = None, None
# 遍历匹配列表，找到最佳匹配
for match, conf, _ in matchList:
# 计算 Levenshtein 距离
d = Levenshtein.distance(text, match)
if best_dist is None or d < best_dist:
best_code, best_dist = match, d
match, conf = best_code, best_dist
# 如果距离大于 2，则认为未识别
# if conf is not None and conf > 2:
# match = "Unidentified"
except IndexError as e:
print(f"An error occurred: {e}")
# 捕获索引越界错误
match = "Unidentified"
except Exception as e:
# 捕获其他异常
print(f"An unexpected error occurred: {e}")
match = "Unidentified"
# 如果匹配结果为 None，则认为未识别
if match is None:
match = "Unidentified"
# 将匹配结果添加到列表中
ol.append(match)
return None, match
def processPoints(points, polygons, items, tag=None):
points = np.array(points, dtype=np.float32)
pg = Polygon(points)
for i, (polygon, t) in enumerate(polygons):
# 计算交集面积
try:
intersection = pg.intersection(polygon)
except GEOSException as e:
print(f"跳过多边形，因为它存在拓扑问题: {e}")
# 可以记录日志或进行其他处理
continue
if not intersection.is_empty:
try:
# 根据面积判断哪个多边形更大
if pg.area > polygon.area:
# pg 减去交集
pg = pg.difference(intersection)
if isinstance(pg, (MultiPolygon, GeometryCollection)):
# 获取所有部分，并找到面积最大的 Polygon
pg = max(pg.geoms, key=lambda p: p.area)
else:
# 当前 polygon 减去交集
polygons[i][0] = polygon.difference(intersection)
if isinstance(polygons[i][0], (MultiPolygon, GeometryCollection)):
# 获取所有部分，并找到面积最大的 Polygon
polygons[i][0] = max(polygons[i][0].geoms, key=lambda p: p.area)
except:
continue
polygons.append([pg, tag])
# items.append([point, center, angle, box, tag, points])
def conduit(conduitList,img_shape):
# import matplotlib.pyplot as plt
conduits, polygons = [], []
for points, tag in conduitList:
processPoints(points, polygons, conduits, tag)
# conduits.append(points)
# fig, ax = plt.subplots(figsize=(10, 10))
for polygon, tag in polygons:
# x, y = polygon.exterior.xy
# ax.plot(x, y, label="Sub-Polygon")
# ax.fill(x, y, alpha=0.4)
coordinates = [[float(x), float(y)] for x, y in list(polygon.exterior.coords)]
rect = cv2.minAreaRect(np.array(coordinates, dtype=np.float32))
center, size, angle = rect
box = cv2.boxPoints(rect) # 获取矩形的四个角点
box = [(item[0], item[1]) for item in np.int32(box)]
boxes = sorted(box, key=lambda box: box[0], reverse=True)
point = [int((boxes[0][0] + boxes[1][0]) / 2), int((boxes[0][1] + boxes[1][1]) / 2)]# 按 x 值排序
conduits.append([point, center, angle, box, tag, coordinates])
conduits.sort(key=lambda x: x[0][1], reverse=True)
conduits = filter_objects_by_line(conduits,img_shape, 500) #进行列聚类
# ax.set_title("title")
# ax.set_xlabel("X Coordinates")
# ax.set_ylabel("Y Coordinates")
# plt.legend()
# plt.axis("equal")
# plt.show()
# print("conduits:", len(conduits), "\tmatches:", len(matches))
return conduits
def nut(clsDict):
nuts, polygons = [], []
if "Nut" not in clsDict: return []
for points in clsDict["Nut"]:
processPoints(points, polygons, nuts)
for polygon, tag in polygons:
coordinates = [[float(x), float(y)] for x, y in list(polygon.exterior.coords)]
rect = cv2.minAreaRect(np.array(coordinates, dtype=np.float32))
center, size, angle = rect
box = cv2.boxPoints(rect) # 获取矩形的四个角点
box = [(item[0], item[1]) for item in np.int32(box)]
boxes = sorted(box, key=lambda box: box[0], reverse=True)
point = [int((boxes[0][0] + boxes[1][0]) / 2), int((boxes[0][1] + boxes[1][1]) / 2)] # 按 x 值排序
nuts.append([point, center, angle, box, tag, coordinates])
nuts.sort(key=lambda x: x[0][1], reverse=True)
return nuts
def find_nearest_coordinate_index(target_coord, coord_array):
"""
找到给定二维坐标在数组中最近的坐标索引。
参数：
target_coord (tuple): 给定的二维坐标 (x, y)。
coord_array (list of tuple): 包含二维坐标的数组 [(x1, y1), (x2, y2), ...]。
返回：
int: 最近坐标的索引值。
"""
if not coord_array:
raise ValueError("坐标数组为空，无法计算最近坐标。")
# 转换为NumPy数组以便向量化计算
coord_array_np = np.array(coord_array)
# 计算欧几里得距离
distances = np.sqrt((coord_array_np[:, 0] - target_coord[0]) ** 2 +
(coord_array_np[:, 1] - target_coord[1]) ** 2)
# 找到最小距离的索引
nearest_index = np.argmin(distances)
return nearest_index
def nose(conduits, clsDict):
if "Nose" not in clsDict: return []
conduitBoxes = [sorted(item[3], key=lambda x: x[0], reverse=True) for item in conduits]
conduitBoxes = [[(box[0][0] + box[1][0]) / 2, (box[0][1] + box[1][1]) / 2] for box in conduitBoxes]
noses, polygons = [], []
for points in clsDict["Nose"]:
processPoints(points, polygons, noses)
for polygon, tag in polygons:
coordinates = [[float(x), float(y)] for x, y in list(polygon.exterior.coords)]
if len(coordinates) == 0: continue
rect = cv2.minAreaRect(np.array(coordinates, dtype=np.float32))
center, size, angle = rect
box = cv2.boxPoints(rect) # 获取矩形的四个角点
box = [(item[0], item[1]) for item in np.int32(box)]
boxes = sorted(box, key=lambda box: box[0], reverse=True)
point = [int((boxes[0][0] + boxes[1][0]) / 2), int((boxes[0][1] + boxes[1][1]) / 2)] # 按 x 值排序
noses.append([point, center, angle, box, tag, coordinates])
noses.sort(key=lambda x: x[0][1], reverse=True)
for i, (_, _, _, box, _, _) in enumerate(noses):
box.sort(key=lambda x: x[0])
leftTop = [(box[0][0] + box[1][0]) / 2, (box[0][1] + box[1][1]) / 2]
belong = find_nearest_coordinate_index(leftTop, conduitBoxes)
noses[i] = [noses[i], belong]
return noses
def progress(img, conduits, noses, nuts, matches,
paras, devicesParas,
response, noseClsModel, detModel, chModel):
def ocrLength(image):
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
ocrRes = chModel.predict(image, verbose=False)[0]
obb = ocrRes.obb.cls.tolist()
return len(obb)
def processConduit(lal, right, cPos, coords, idx):
OK[idx][1] = 0
occ = " ".join([lal, f"#{cPos + 1:02d}", right])
if occ in occur: return
occur.append(occ)
insertResponse(device, lal, coords, f"#{cPos + 1:02d}", right, response)
def checkConduit(mch, idx):
if mch in halfDevices:
for j, item in enumerate(corrects):
if item[0] == mch:
if matches[j][0] == corrects[idx][0] or matches[j][0] == "Unidentified":
_, _, _, _, _, pts, _ = conduits[j]
pts = inverseTransformCoordinates(pts, rotate, img.shape)
processConduit("EXP_CONNECTOR_LINE_POSITION", c, cp, pts, j)
return False
return True
else:
return True
outList = []
features, device, rotate = devicesParas
devices = features[device]["set"]
deviceLenght = len(devices) / 2
conduitBoxes = [sorted(item[3], key=lambda x: x[0], reverse=True) for item in conduits]
conduitBoxes = [[(box[0][0] + box[1][0]) / 2, (box[0][1] + box[1][1]) / 2] for box in conduitBoxes]
if deviceLenght < len(conduits):
xList = [item[3][0][0] for item in nuts]
if len(xList) < len(conduits):
for i in range(len(conduits) - len(xList)):
xList.append(xList[-1])
popLength = len(conduits) - deviceLenght
while popLength > 0:
conduitsXList = [item[3][1][0] for item in conduits]
dist = [nut - conduit for conduit, nut in zip(conduitsXList, xList)]
maxDist = dist.index(max(dist))
outList.append([conduits.pop(maxDist), maxDist])
matches.pop(maxDist)
popLength -= 1
rev, counter, pos = paras
corrects = copy.deepcopy(matches)
for i, item in enumerate(corrects):
offset = (i - pos) if rev else pos - i
p = counter + offset
if p < len(devices):
corrects[i] = [devices[p], p]
else:
corrects[i] = ["ERR", p]
OK = [[correct[0], 1] for correct in corrects]
halfDevices = [it[0] for it in OK]
for i in range(len(corrects)):
# if matches[i][0] != "Unidentified": continue
# matches[i] = corrects[i]
print(f"#{corrects[i][1] + 1}: ", matches[i][0], "\t------->", corrects[i][0])
for out in outList:
for match, correct in zip(matches, corrects):
if match[0] == "empty" and correct[0] != "empty": out.append(correct[1])
occur = []
for i, cdu in enumerate(conduits):
point, center, angle, box, tag, points, crop = cdu
m, mp = matches[i]
c, cp = corrects[i]
points = inverseTransformCoordinates(points, rotate, img.shape)
if m == "Unidentified" and c == "empty":
processConduit("EXP_CONNECTOR_LINE_POSITION", c, cp, points, i)
elif m == "Unidentified" and c != "empty":
processConduit("WARN_CONFIRM", c, cp, points, i)
elif m == "empty" and c != "empty":
processConduit("EXP_CONNECTOR_LINE_EMPTY", c, cp, points, i)
elif m != c and m not in halfDevices:
processConduit("WARN_CONFIRM", c, cp, points, i)
elif m != c and m in halfDevices:
lenMatch = ocrLength(crop)
if lenMatch > len(m):
processConduit("WARN_CONFIRM", c, cp, points, i)
continue
processConduit("EXP_CONNECTOR_LINE_POSITION", c, cp, points, i)
unchange = checkConduit(m, i)
if not unchange: continue
response["result"].pop(-1)
occur.pop(-1)
processConduit("WARN_CONFIRM", c, cp, points, i)
# if len(noses) > 0:
# for (point, center, angle, box, tag, points), belong in noses:
# mask = np.zeros(img.shape[:2], dtype=np.uint8)
# # 填充多边形区域为白色
# cv2.fillPoly(mask, [np.array(points, dtype=np.int32)], 255)
# res = cv2.bitwise_and(img, img, mask=mask)
# if belong > deviceLenght - 1: continue
# correct = corrects[belong]
# points = inverseTransformCoordinates(box, rotate, img.shape)
# res = cv2.bitwise_and(img, img, mask=mask)
# gray = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY) # 转为灰度图
# _, binary = cv2.threshold(gray, 3, 255, cv2.THRESH_BINARY) # 将非黑色部分设为白色
# # 计算非黑色区域的边界框
# x, y, w, h = cv2.boundingRect(binary)
# if w > 0 and h > 0:
# # 裁剪非黑色部分
# res = res[y:y + h, x:x + w]
# status = image2Cls(res, noseClsModel)
# if status == "ls":
# OK[belong][1] = 0
# occ = " ".join(["EXP_CONNECTOR_LINE_UNSEATED", f"#{correct[1] + 1:02d}", correct[0]])
# if occ in occur: continue
# occur.append(occ)
# insertResponse(device, "EXP_CONNECTOR_LINE_UNSEATED", points, f"#{correct[1] + 1:02d}", correct[0], response)
ress = detModel.predict(img, save=True, verbose=False, exist_ok=True)
names = list(detModel.names.values())
if ress[0] is not None:
for res in ress[0]:
clss = [int(item) for item in res.boxes.cls.tolist()]
boxes = res.boxes.xyxy.tolist()
for c, b in zip(clss, boxes):
label = names[c]
x_min, y_min, x_max, y_max = b
# 矩形的四个顶点
box = [
[x_min, y_min], # 左上角
[x_max, y_min], # 右上角
[x_max, y_max], # 右下角
[x_min, y_max], # 左下角
]
if label == "Unconnected":
belong = find_nearest_coordinate_index([x_min, (y_min + y_max) / 2], conduitBoxes)
points = inverseTransformCoordinates(box, rotate, img.shape)
if belong > deviceLenght - 1:
for out in outList:
if int(out[1]) == int(belong):
OK[out[2]][1] = 0
occ = " ".join(["EXP_CONNECTOR_LINE_UNSEATED", f"#{out[2] + 1:02d}", devices[out[2]]])
if occ in occur: continue
occur.append(occ)
insertResponse(device, "EXP_CONNECTOR_LINE_UNSEATED", points, f"#{out[2] + 1:02d}",
devices[out[2]], response)
else:
correct = corrects[belong]
OK[belong][1] = 0
occ = " ".join(["EXP_CONNECTOR_LINE_UNSEATED", f"#{correct[1] + 1:02d}", correct[0]])
if occ in occur: continue
occur.append(occ)
insertResponse(device, "EXP_CONNECTOR_LINE_UNSEATED", points, f"#{correct[1] + 1:02d}",
correct[0], response)
# TODO
for i, ok in enumerate(OK):
if ok[1] == 1:
correct = corrects[i]
point, center, angle, box, tag, points, _ = conduits[i]
points = inverseTransformCoordinates(box, rotate, img.shape)
insertResponse(device, "OK", points, f"#{correct[1] + 1:02d}", correct[0], response)
print("done")
def insertResponse(device, tp, roi, port, lan, response):
response["result"].append({
"device": device,
"type": tp,
"roi": roi,
"port": port,
"lan": lan,
})
def drawResults(response, image, rotate):
img = cv2.imread(image)
org = img.shape[:2]
scale = 3 if org[1] > 2000 else 1
org = (org[1] - 400, 100)
if response is None:
cv2.putText(img, "WARN", org, cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 255, 255), 5)
else:
results = response["result"]
flag = 0
# 存储已绘制的矩形区域
rectangles = []
for result in results:
roi = result["roi"]
polygon = Polygon(roi)
box = list(polygon.minimum_rotated_rectangle.exterior.coords)
contour = np.array(box, dtype=np.int32).reshape((-1, 1, 2))
if result["type"] == "OK":
color = (0, 255, 0)
cv2.drawContours(img, [contour], -1, color, 3)
elif result["type"] == "WARN_CONFIRM":
color = (0, 255, 255)
flag += 1
print(result["type"], result["port"], result["lan"])
cv2.drawContours(img, [contour], -1, color, 3)
drawText(img, box, scale, result["lan"], rectangles, color, rotate)
elif result["type"] == "EXP_CONNECTOR_LINE_POSITION":
color = (0, 0, 255)
flag += 1
print(result["type"], result["port"],result["lan"])
cv2.drawContours(img, [contour], -1, color, 3)
drawText(img, box, scale, result["lan"], rectangles, color, rotate)
elif result["type"] == "EXP_CONNECTOR_LINE_EMPTY":
color = (0, 0, 255)
flag += 1
print(result["type"], result["port"],result["lan"])
cv2.drawContours(img, [contour], -1, color, 3)
drawText(img, box, scale, result["lan"], rectangles, color, rotate)
else:
color = (0, 0, 255)
flag += 1
print(result["type"], result["port"], result["lan"])
cv2.drawContours(img, [contour], -1, color, 3)
drawText(img, box, scale, result["port"] + " " + result["type"], rectangles, color, rotate)
if flag == 0:
cv2.putText(img, "OK", org, cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 255, 0), 5)
if img.shape[1] > 2000:
img = cv2.resize(img, dsize=(0, 0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA)
elif 2000 > img.shape[1] > 1000:
img = cv2.resize(img, dsize=(0, 0), fx=0.6, fy=0.6, interpolation=cv2.INTER_AREA)
# cv2.imwrite(output, img)
# cv2.imshow("result", img)
# cv2.waitKey(0)
return img
def drawText(img, box, scale, label, rectangles, bgColor, rotate):
coord = min(box, key=lambda x: x[1]) if rotate < 180 else max(box, key=lambda x: x[1])
coord = (int(coord[0]), int(coord[1]))
cv2.circle(img, (int(coord[0]), int(coord[1])), 5 * scale, bgColor, -1)
if scale < 2:
if int(rotate) < 180:
top_left = [int(coord[0]) + 20 * scale,
int(coord[1] - 105)]
bottom_right = [int(coord[0]) + 50 + 12 * len(label) * scale,
int(coord[1] - 65)]
else:
top_left = [int(coord[0]) + 20 * scale,
int(coord[1] + 65)]
bottom_right = [int(coord[0]) + 50 + 12 * len(label) * scale,
int(coord[1] + 105)]
else:
if int(rotate) < 180:
top_left = [int(coord[0]) + 20 * scale,
int(coord[1] - 145)]
bottom_right = [int(coord[0]) + 130 + 12 * len(label) * scale,
int(coord[1] - 65)]
else:
top_left = [int(coord[0]) + 20 * scale,
int(coord[1] + 65)]
bottom_right = [int(coord[0]) + 130 + 12 * len(label) * scale,
int(coord[1] + 145)]
# 确保矩形不与其他已绘制矩形重叠
while any(
not (
top_left[0] > rect[1][0] or # 当前矩形在其他矩形的右边
bottom_right[0] < rect[0][0] or # 当前矩形在其他矩形的左边
top_left[1] > rect[1][1] or # 当前矩形在其他矩形的下边
bottom_right[1] < rect[0][1] # 当前矩形在其他矩形的上边
)
for rect in rectangles
):
if rotate < 180:
top_left[1] -= 20 # 向上移动
bottom_right[1] -= 20
else:
top_left[1] += 20 # 向下移动
bottom_right[1] += 20
# 记录当前矩形
rectangles.append((top_left, bottom_right))
labelPos = (top_left[0], bottom_right[1])
cv2.line(img, coord, labelPos, bgColor, 2)
# 绘制矩形和标签
cv2.rectangle(img, tuple(top_left), tuple(bottom_right), bgColor, -1)
# 绘制文字
font = cv2.FONT_HERSHEY_SIMPLEX # 字体类型
font_scale = 0.6 * scale # 字体大小
font_thickness = 1 * scale # 字体线条厚度
text_color = (0, 0, 0) # 黑色文字
text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
text_width, text_height = text_size
rect_center_x = (top_left[0] + bottom_right[0]) // 2
rect_center_y = (top_left[1] + bottom_right[1]) // 2
# 计算文字的左上角坐标（使文字中心与矩形中心对齐）
text_x = rect_center_x - text_width // 2
text_y = rect_center_y + text_height // 2
cv2.putText(img, label, (text_x, text_y), font, font_scale, text_color, font_thickness)
def terminal_process(path, point):
SegModel = YOLO("models/terminal_seg_7_15_v3.pt")
orientedClsModel = YOLO("models/orientedCls.pt")
noseClsModel = YOLO("models/noseCls.pt")
detModel = YOLO("models/det.pt")
chModel = YOLO("models/ch.pt")
outarea = YOLO("models/outarea.pt")
features = praseExcel("dataset/excel/logbook-0.xlsx")
index = 0
response = {
"result": []
}
if path.endswith((".jpg", ".jpeg", ".png")):
im = cv2.imread(path)
img, textImg, clsDict, results, rotate = image2Align(path, orientedClsModel, SegModel, noseClsModel)
img_shape = img.shape
# cv2.imwrite(os.path.join(path.replace(".jpg", "_0.jpg")), img)
emptyList = [[item, "empty"] for item in clsDict["Empty"]] if "Empty" in clsDict else []
conduitList = [[item, "Conduit"] for item in clsDict["Conduit"]] if "Conduit" in clsDict else []
nuts = nut(clsDict)
conduits = conduit(emptyList + conduitList,img_shape)
noses = nose(conduits, clsDict)
matches, device, paras = image2OCR(img, conduits, features, os.path.basename(path))
if matches is not None:
progress(img, conduits, noses, nuts, matches,
paras, [features, device, rotate],
response, noseClsModel, detModel, chModel)
else:
return None, None
return response, rotate
def run(path,output_folder,SegModel,orientedClsModel,noseClsModel,detModel,chModel,features,index):
for file in sorted(os.listdir(path)):
if file.endswith((".jpg", ".jpeg", ".png")):
p = os.path.join(path, file)
# p_shape = cv2.imread(p).shape
o = os.path.join(output_folder, file)
# p = "//home/master/KR0418/12-PingGui/result/raw/26.jpg"
print(p)
r, rot = terminal_process(p, "", SegModel, orientedClsModel, noseClsModel, detModel, chModel, features, index)
drawResults(r, p, o, rot)
# print(f"\n本次测试已结束，测试图片共{len(os.listdir(path))}张，错误图片共{wrong_pic_count}张，测试准确率为：{(len(os.listdir(path))-wrong_pic_count)/len(os.listdir(path))}\n")
if __name__ == '__main__':
image_filepath = '/home/pbl/code/KR0530/12-PingGui_v5/12-PingGui/04b03bafc7984fdb93190dadb59a039e.jpg'
r, rot = terminal_process(image_filepath, '')
# print(r) EXP_CONNECTOR_LINE_EMPTY
det_img= drawResults(r, image_filepath, rot)
cv2.imwrite('/home/pbl/code/KR0530/12-PingGui_v5/12-PingGui/output1.jpg', det_img)
在这个程序中，被检测出有错误的地方是怎么被画成红框的，是怎么在图像上显示错误的
最新发布
本文介绍了一种在已知有序数组被旋转后的场景下寻找数组最小值的有效算法。该算法利用二分查找的方式，通过比较数组两端及中间元素的大小来逐步缩小搜索范围，最终定位到最小值。
扫一扫
228
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
