DeepStream--测试PCB-Defect-Detection

GitHub - clintonoduor/PCB-Defect-Detection-using-Deepstream: PCB defect detection using deepstream & YoloV5我参考了了这个代码，作者基于YoloV5，训练一个电路板检测的模型，训练数据集来自https://robotics.pkusz.edu.cn/resources/datasetENG/。模型支持以下6种残缺的检测：

Missing hole 漏孔
Mouse bite 鼠牙洞
Open circuit 开路
Short circuit 短路
Spurious copper 杂铜
spur 毛刺

但是这个代码有个问题，工程缺少custom_yolov5s.yaml 文件，这一步将失败，生成不了wts和cfg文件。

!python3 gen_wts_yoloV5.py -w /content/yolov5/runs/train/yolov5s_results/weights/best.pt -c /content/yolov5/models/custom_yolov5s.yaml

我按https://github.com/marcoslucianops/DeepStream-Yolo/blob/master/docs/YOLOv5.md 这里的步骤，将pt文件直接转成onnx模型。最终转的命令是：python3 export_yoloV5.py -w best.pt

config_infer_primary_yoloV5.txt的内容：

[property]
gpu-id=0
net-scale-factor=0.0039215697906911373
model-color-format=0
onnx-file=best.onnx
model-engine-file=best.onnx_b1_gpu0_fp32.engine
#int8-calib-file=calib.table
labelfile-path=labels.txt
batch-size=1
network-mode=0
num-detected-classes=6
interval=0
gie-unique-id=1
process-mode=1
network-type=0
cluster-mode=4
maintain-aspect-ratio=1
parse-bbox-func-name=NvDsInferParseYolo
custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so
#engine-create-func-name=NvDsInferYoloCudaEngineGet

[class-attrs-all]
pre-cluster-threshold=0.2

后处理函数用NvDsInferParseCustomYolo。

static NvDsInferParseObjectInfo
convertBBox(const float& bx1, const float& by1, const float& bx2, const float& by2, const uint& netW, const uint& netH)
{NvDsInferParseObjectInfo b;float x1 = bx1;float y1 = by1;float x2 = bx2;float y2 = by2;x1 = clamp(x1, 0, netW);y1 = clamp(y1, 0, netH);x2 = clamp(x2, 0, netW);y2 = clamp(y2, 0, netH);b.left = x1;b.width = clamp(x2 - x1, 0, netW);b.top = y1;b.height = clamp(y2 - y1, 0, netH);return b;
}static void
addBBoxProposal(const float bx1, const float by1, const float bx2, const float by2, const uint& netW, const uint& netH,const int maxIndex, const float maxProb, std::vector<NvDsInferParseObjectInfo>& binfo)
{NvDsInferParseObjectInfo bbi = convertBBox(bx1, by1, bx2, by2, netW, netH);if (bbi.width < 1 || bbi.height < 1)return;bbi.detectionConfidence = maxProb;bbi.classId = maxIndex;binfo.push_back(bbi);
}static std::vector<NvDsInferParseObjectInfo>
decodeTensorYolo(const float* boxes, const float* scores, const float* classes, const uint& outputSize, const uint& netW,const uint& netH, const std::vector<float>& preclusterThreshold)
{std::vector<NvDsInferParseObjectInfo> binfo;for (uint b = 0; b < outputSize; ++b) {float maxProb = scores[b];int maxIndex = (int) classes[b];if (maxProb < preclusterThreshold[maxIndex])continue;float bxc = boxes[b * 4 + 0];float byc = boxes[b * 4 + 1];float bw = boxes[b * 4 + 2];float bh = boxes[b * 4 + 3];float bx1 = bxc - bw / 2;float by1 = byc - bh / 2;float bx2 = bx1 + bw;float by2 = by1 + bh;addBBoxProposal(bx1, by1, bx2, by2, netW, netH, maxIndex, maxProb, binfo);}return binfo;
}static bool
NvDsInferParseCustomYolo(std::vector<NvDsInferLayerInfo> const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo,NvDsInferParseDetectionParams const& detectionParams, std::vector<NvDsInferParseObjectInfo>& objectList)
{if (outputLayersInfo.empty()) {std::cerr << "ERROR: Could not find output layer in bbox parsing" << std::endl;return false;}std::vector<NvDsInferParseObjectInfo> objects;const NvDsInferLayerInfo& boxes = outputLayersInfo[0];const NvDsInferLayerInfo& scores = outputLayersInfo[1];const NvDsInferLayerInfo& classes = outputLayersInfo[2];const uint outputSize = boxes.inferDims.d[0];std::vector<NvDsInferParseObjectInfo> outObjs = decodeTensorYolo((const float*) (boxes.buffer),(const float*) (scores.buffer), (const float*) (classes.buffer), outputSize, networkInfo.width, networkInfo.height,detectionParams.perClassPreclusterThreshold);objects.insert(objects.end(), outObjs.begin(), outObjs.end());objectList = objects;return true;
}

测试图片：

测试命令：

gst-launch-1.0 filesrc location=a.JPG ! jpegdec ! videoconvert ! video/x-raw,format=I420 ! nvvideoconvert ! video/x-raw\(memory:NVMM\),format=NV12 ! mux.sink_0 nvstreammux name=mux batch-size=1 width=1280 height=720 ! nvinfer config-file-path=./config_infer_primary_yoloV5.txt ! nvvideoconvert ! video/x-raw\(memory:NVMM\),format=RGBA ! nvdsosd ! nvvideoconvert ! video/x-raw,format=I420 ! jpegenc ! filesink location=out.jpg

结果图片：