使用YOLOv8进行工业视觉缺陷检测，基于Aidlux完成本地终端部署

class VoVGSCSP(nn.Module):# VoVGSCSP module with GSBottleneckdef __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):super().__init__()c_ = int(c2 * e)  # hidden channelsself.cv1 = Conv(c1, c_, 1, 1)self.cv2 = Conv(c1, c_, 1, 1)self.gsb = nn.Sequential(*(GSBottleneck(c_, c_, e=1.0) for _ in range(n)))self.res = Conv(c_, c_, 3, 1, act=False)self.cv3 = Conv(2 * c_, c2, 1)def forward(self, x):x1 = self.gsb(self.cv1(x))y = self.cv2(x)return self.cv3(torch.cat((y, x1), dim=1))class VoVGSCSPC(VoVGSCSP):# cheap VoVGSCSP module with GSBottleneckdef __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):super().__init__(c1, c2)c_ = int(c2 * 0.5)  # hidden channelsself.gsb = GSBottleneckC(c_, c_, 1, 1)

    if mode=="onnx" :model = YOLO('D:\\gongye\\yolov8-main\\gongye.pt')model.export(format="onnx",opset=11,simplify=True)# path = model.export(format="onnx",opset=13,half=True,simplify=True)# path = model.export(format="onnx",opset=13,half=True,simplify=True)

import aidlite_gpu
import cv2
from cvs import *
import numpy as np
import os
import timemodel_path = "/home/gongye/tflite/gongye.tflite"
image_path = "/home/gongye/test"
NUMS_CLASS = 7confThresh = 0.3
NmsThresh = 0.45# 输入格式 （8400,11）
def postProcess(pres, confThresh, NmsThresh):boxes_out = []scores_out = []class_out = []for pred in pres:pred_class = pred[4:]box_ = pred[0:4]# pred_class=(pred_class-min(pred_class))/(max(pred_class)-min(pred_class))class_index = np.argmax(pred_class)if pred_class[class_index] > 0.3:# box=np.array([round(pred[2]-0.5*pred[0]),round(pred[3]-0.5*pred[1]),round(pred[0]),round(pred[1])])box_ = pred[0:4]  # w,h,xc,ycbox = np.array([round((pred[2] / 2 - pred[0])), round((pred[3] / 2 - pred[1])), round(pred[0] * 2),round(pred[1] * 2)])boxes_out.append(box)score = pred_class[class_index]scores_out.append(score)class_out.append(class_index)result_boxes = cv2.dnn.NMSBoxes(boxes_out, np.array(scores_out), confThresh, NmsThresh)# detections=[]boxes = []scores = []classes = []for result_box in result_boxes:index = int(result_box)box = boxes_out[index]score = scores_out[index]class_type = class_out[index]boxes.append(box)scores.append(score)classes.append(class_type)return boxes, scores, classesdef draw(img, xscale, yscale, boxes, scores, classes):width = img.shape[1]w1 = 1620w2 = 2350w3 = widthS1 = []S2 = []S3 = []S1_res = [False for i in range(NUMS_CLASS)]S2_res = [False for i in range(NUMS_CLASS)]S3_res = [False for i in range(NUMS_CLASS)]S_res = [S1_res, S2_res, S3_res]img_ = img.copy()# 遍历所有box，按照分割区域将box归类for i in range(len(boxes)):# boxes=[x1,y1,w,h]box = boxes[i]score = scores[i]class_ = int(classes[i])# class_text=label[class_]# detect=[round(box[0]*xscale),round(box[1]*yscale),round((box[0]+box[2])*xscale),round((box[1]+box[3])*yscale)]detect = [round(box[0] * xscale), round(box[1] * yscale), round(box[0] * xscale + (box[2]) * xscale),round(box[1] * yscale + (box[3]) * yscale)]text = "{}:{:.2f}".format(label[class_], float(score))img_ = cv2.rectangle(img_, (detect[0], detect[1]), (detect[2], detect[3]), (0, 255, 0), 2)cv2.putText(img_, text, (detect[0], detect[1] + 10), cv2.FONT_HERSHEY_COMPLEX, 2, (0, 0, 255), 1)# 分割为三块if (detect[0] <= w1):p1 = []p1.append(detect)p1.append(class_)p1.append(score)S1.append(p1)elif (w1 < detect[0] <= w2):p2 = []p2.append(detect)p2.append(class_)p2.append(score)S2.append(p2)elif (w2 < detect[0] <= w3):p3 = []p3.append(detect)p3.append(class_)p3.append(score)S3.append(p3)# 判断每个分割图像中的结果index = 0for S in [S1, S2, S3]:for i in range(len(S)):p1 = S[i]box_temp = p1[0]class_temp = p1[1]score_temp = p1[2]S_res[index][class_temp] = Trueindex += 1# 最终分割输出结果true or falseS_out = [False, False, False]index_out = 0for s_r in S_res:c0 = s_r[0]c1 = s_r[1]c2 = s_r[2]c3 = s_r[3]c4 = s_r[4]c5 = s_r[5]c6 = s_r[6]if (c0 & c1 & c2 & c3 & (~c4) & (~c5) & (~c6)):S_out[index_out] = Trueelif (c0 & c1 & c2 & (~c3) & (~c4) & c5 & (~c6)):S_out[index_out] = Trueindex_out += 1# 打印分割结果cv2.putText(img_, "OK" if S_out[0] == True else "NG", (w1 - 200, 100), cv2.FONT_HERSHEY_COMPLEX, 2, (255, 0, 0), 1)cv2.putText(img_, "OK" if S_out[1] == True else "NG", (w2 - 200, 100), cv2.FONT_HERSHEY_COMPLEX, 2, (255, 0, 0), 1)cv2.putText(img_, "OK" if S_out[2] == True else "NG", (w3 - 200, 100), cv2.FONT_HERSHEY_COMPLEX, 2, (255, 0, 0), 1)return img_label = ["rubber stopper", "push rod tail", "needle tail", "mouth", "crooked mouth", "screw mouth", "small rubber plug"]if __name__ == "__main__":# 1.初始化aidlite类并创建aidlite对象aidlite = aidlite_gpu.aidlite()print("ok")# 2.加载模型value = aidlite.ANNModel(model_path, [640 * 640 * 3 * 4], [8400 * 11 * 4], 4, 0)print("gpu:", value)# file_names=os.listdir(image_path)# root,dirs,files = os.walk(image_path)for root, dirs, files in os.walk(image_path):num = 0for file in files:file = os.path.join(root, file)frame = cv2.imread(file)x_scale = frame.shape[1] / 640y_scale = frame.shape[0] / 640img = cv2.resize(frame, (640, 640))# img_copy=img.coimg = img / 255.0img = np.expand_dims(img, axis=0)img = img.astype(dtype=np.float32)print(img.shape)# 3.传入模型输入数据aidlite.setInput_Float32(img)# 4.执行推理start = time.time()aidlite.invoke()end = time.time()timerValue = 1000 * (end - start)print("infer time(ms):{0}", timerValue)# 5.获取输出pred = aidlite.getOutput_Float32(0)# print(pred.shape)pred = np.array(pred)print(pred.shape)pred = np.reshape(pred, (8400, 11))# pred=np.reshape(pred,(11,8400)).transpose()print(pred.shape)  # shape=(8400,11)# 6.后处理,解析输出boxes, scores, classes = postProcess(pred, confThresh, NmsThresh)# 7.绘制保存图像ret_img = draw(frame, x_scale, y_scale, boxes, scores, classes)ret_img = ret_img[:, :, ::-1]num += 1image_file_name = "/home/gongye/result/tfl_res" + str(num) + ".jpg"# 8.保存图片cv2.imwrite(image_file_name, ret_img)