Boen_Shi
6a2e046884
- 新增 analyze、analyze_result、analyze_status 和 health 路由 - 实现图像上传和任务提交功能 - 添加任务状态查询和结果获取接口 - 集成 segformer 和 yolo 模型进行图像检测 - 实现 SAM3 预处理功能用于图像预处理判断 - 添加模型选择配置支持 segformer 和 yolo - 实现任务队列管理和异步处理机制 - 添加 Dockerfile 用于容器化部署 - 配置环境变量和 gitignore 规则 - 创建数据模型定义 API 响应结构
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import albumentations as A
|
|
import onnxruntime as ort
|
|
|
|
|
|
class Detection:
|
|
def __init__(self):
|
|
self.CLASS_NAMES = ["background", "Hollowing", "Water seepage", "Cracking"]
|
|
self.PALETTE = np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]], dtype=np.uint8)
|
|
self.tfm = A.Compose([
|
|
A.Resize(512, 512),
|
|
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
real_path = os.path.join(os.path.dirname(__file__), "segformer.onnx")
|
|
self.model = ort.InferenceSession(
|
|
real_path,
|
|
providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
)
|
|
print("ONNX 模型加载完成！")
|
|
|
|
def get_contours(self, mask_np):
|
|
res = []
|
|
for idx in range(1, len(self.CLASS_NAMES)):
|
|
name = self.CLASS_NAMES[idx]
|
|
binary = (mask_np == idx).astype(np.uint8) * 255
|
|
cnts, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
for c in cnts:
|
|
if cv2.contourArea(c) < 20:
|
|
continue
|
|
eps = 0.002 * cv2.arcLength(c, True)
|
|
poly = cv2.approxPolyDP(c, eps, True)
|
|
if len(poly) >= 3:
|
|
res.append((name, poly.reshape(-1, 2).tolist()))
|
|
return res
|
|
|
|
def detect(self, img_input):
|
|
# 读取图片
|
|
if isinstance(img_input, str):
|
|
img_bgr = cv2.imdecode(np.fromfile(img_input, dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
else:
|
|
img_bgr = img_input
|
|
|
|
h, w = img_bgr.shape[:2]
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
# 数据预处理
|
|
x = self.tfm(image=img_rgb)["image"]
|
|
# 将 HWC -> CHW 并增加 batch 维度
|
|
x = np.transpose(x, (2, 0, 1)) # HWC -> CHW
|
|
x = np.expand_dims(x, axis=0).astype(np.float32) # batch x C x H x W
|
|
|
|
# ONNX 推理
|
|
inp_name = self.model.get_inputs()[0].name
|
|
out_name = self.model.get_outputs()[0].name
|
|
out = self.model.run([out_name], {inp_name: x})[0]
|
|
|
|
# 获取预测结果
|
|
pred = out.argmax(axis=1)[0].astype(np.uint8)
|
|
|
|
mask = cv2.resize(pred, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
mask_rgb = self.PALETTE[mask]
|
|
coords = self.get_contours(mask)
|
|
# print(coords)
|
|
return mask_rgb, coords
|
|
|
|
|
|
class DetectionMock:
|
|
def __init__(self):
|
|
self.onnx_path = "segformer.onnx"
|
|
|
|
def detect(self, img_input):
|
|
if isinstance(img_input, str):
|
|
img_bgr = cv2.imdecode(np.fromfile(img_input, dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
else:
|
|
img_bgr = img_input
|
|
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
return img_rgb, [
|
|
('Cracking', [[771, 928], [757, 935]]),
|
|
('Cracking', [[254, 740], [251, 942]]),
|
|
('Cracking', [[764, 420], [764, 424]]),
|
|
('Cracking', [[257, 238], [251, 245]]),
|
|
('Cracking', [[1436, 145], [1401, 145]])
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
img_path = "test.jpg"
|
|
|
|
detection = Detection()
|
|
mask_res, coords_res = detection.detect(img_path)
|
|
|
|
print("Mask Shape:", mask_res.shape)
|
|
print("Detections:", coords_res)
|
|
cv2.imwrite("res_onnx.png", cv2.cvtColor(mask_res, cv2.COLOR_RGB2BGR))
|