YOLOV----- ONNX 推理过程、可视化图片、保存检测到的目标
创始人
2024-09-25 06:22:59
0

一、代码

import os import cv2 import numpy as np import onnxruntime import time  CLASSES = ['person']  # coco80类别   class YOLOV5():     def __init__(self, onnxpath):         self.onnx_session = onnxruntime.InferenceSession(onnxpath)         self.input_name = self.get_input_name()         self.output_name = self.get_output_name()      # -------------------------------------------------------     #   获取输入输出的名字     # -------------------------------------------------------     def get_input_name(self):         input_name = []         for node in self.onnx_session.get_inputs():             input_name.append(node.name)         return input_name      def get_output_name(self):         output_name = []         for node in self.onnx_session.get_outputs():             output_name.append(node.name)         return output_name      # -------------------------------------------------------     #   输入图像     # -------------------------------------------------------     def get_input_feed(self, img_tensor):         input_feed = {}         for name in self.input_name:             input_feed[name] = img_tensor         return input_feed      # -------------------------------------------------------     #   1.cv2读取图像并resize     #	2.图像转BGR2RGB和HWC2CHW     #	3.图像归一化     #	4.图像增加维度     #	5.onnx_session 推理     # -------------------------------------------------------     # def inference(self, img_path):     #     img = cv2.imread(img_path)     #     or_img = cv2.resize(img, (640, 640))     #     img = or_img[:, :, ::-1].transpose(2, 0, 1)  # BGR2RGB和HWC2CHW     #     img = img.astype(dtype=np.float32)     #     img /= 255.0     #     img = np.expand_dims(img, axis=0)     #     input_feed = self.get_input_feed(img)     #     pred = self.onnx_session.run(None, input_feed)[0]     #     return pred, or_img      # def inference(self, img_path):     #     try:     #         img = cv2.imread(img_path)     #         or_img = cv2.resize(img, (640, 640))     #         img = or_img[:, :, ::-1].transpose(2, 0, 1)  # BGR2RGB和HWC2CHW     #         img = img.astype(dtype=np.float32) / 255.0     #         img = np.expand_dims(img, axis=0)     #         input_feed = self.get_input_feed(img)     #         pred = self.onnx_session.run(None, input_feed)[0]     #         return pred, or_img     #     except Exception as e:     #         print(f"Error during inference: {e}")     #         return None, None      def inference(self, img):         try:             img = cv2.imread(img)             or_img = img.copy()             resized_img = cv2.resize(img, (640, 640))  # 根据模型要求进行resize             resized_img = resized_img[:, :, ::-1].transpose(2, 0, 1)  # BGR2RGB和HWC2CHW             resized_img = resized_img.astype(dtype=np.float32) / 255.0             resized_img = np.expand_dims(resized_img, axis=0)             input_feed = self.get_input_feed(resized_img)             pred = self.onnx_session.run(None, input_feed)[0]             return pred, or_img, resized_img         except Exception as e:             print(f"Error during inference: {e}")             return None, None, None   # dets:  array [x,6] 6个值分别为x1,y1,x2,y2,score,class # thresh: 阈值 def nms(dets, thresh):     x1 = dets[:, 0]     y1 = dets[:, 1]     x2 = dets[:, 2]     y2 = dets[:, 3]     # -------------------------------------------------------     #   计算框的面积     #	置信度从大到小排序     # -------------------------------------------------------     areas = (y2 - y1 + 1) * (x2 - x1 + 1)     scores = dets[:, 4]     keep = []     index = scores.argsort()[::-1]      while index.size > 0:         i = index[0]         keep.append(i)         # -------------------------------------------------------         #   计算相交面积         #	1.相交         #	2.不相交         # -------------------------------------------------------         x11 = np.maximum(x1[i], x1[index[1:]])         y11 = np.maximum(y1[i], y1[index[1:]])         x22 = np.minimum(x2[i], x2[index[1:]])         y22 = np.minimum(y2[i], y2[index[1:]])          w = np.maximum(0, x22 - x11 + 1)         h = np.maximum(0, y22 - y11 + 1)          overlaps = w * h         # -------------------------------------------------------         #   计算该框与其它框的IOU,去除掉重复的框,即IOU值大的框         #	IOU小于thresh的框保留下来         # -------------------------------------------------------         ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)         idx = np.where(ious <= thresh)[0]         index = index[idx + 1]     return keep   def xywh2xyxy(x):     # [x, y, w, h] to [x1, y1, x2, y2]     y = np.copy(x)     y[:, 0] = x[:, 0] - x[:, 2] / 2     y[:, 1] = x[:, 1] - x[:, 3] / 2     y[:, 2] = x[:, 0] + x[:, 2] / 2     y[:, 3] = x[:, 1] + x[:, 3] / 2     return y   def filter_box(org_box, conf_thres, iou_thres):  # 过滤掉无用的框     # -------------------------------------------------------     #   删除为1的维度     #	删除置信度小于conf_thres的BOX     # -------------------------------------------------------     org_box = np.squeeze(org_box)     conf = org_box[..., 4] > conf_thres     box = org_box[conf == True]     # -------------------------------------------------------     #	通过argmax获取置信度最大的类别     # -------------------------------------------------------     cls_cinf = box[..., 5:]     cls = []     for i in range(len(cls_cinf)):         cls.append(int(np.argmax(cls_cinf[i])))     all_cls = list(set(cls))     # -------------------------------------------------------     #   分别对每个类别进行过滤     #	1.将第6列元素替换为类别下标     #	2.xywh2xyxy 坐标转换     #	3.经过非极大抑制后输出的BOX下标     #	4.利用下标取出非极大抑制后的BOX     # -------------------------------------------------------     output = []     for i in range(len(all_cls)):         curr_cls = all_cls[i]         curr_cls_box = []         curr_out_box = []         for j in range(len(cls)):             if cls[j] == curr_cls:                 box[j][5] = curr_cls                 curr_cls_box.append(box[j][:6])         curr_cls_box = np.array(curr_cls_box)         # curr_cls_box_old = np.copy(curr_cls_box)         curr_cls_box = xywh2xyxy(curr_cls_box)         curr_out_box = nms(curr_cls_box, iou_thres)         for k in curr_out_box:             output.append(curr_cls_box[k])     output = np.array(output)     return output   def draw(image, box_data):     # -------------------------------------------------------     #	取整,方便画框     # -------------------------------------------------------     boxes = box_data[..., :4].astype(np.int32)     scores = box_data[..., 4]     classes = box_data[..., 5].astype(np.int32)      for box, score, cl in zip(boxes, scores, classes):         top, left, right, bottom = box         print('class: {}, score: {}'.format(CLASSES[cl], score))         print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))          cv2.rectangle(image, (top, left), (right, bottom), (0, 0, 255), 2) #红色         #cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2) #蓝色         cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),                     (top, left),                     cv2.FONT_HERSHEY_SIMPLEX,                     0.6, (0, 0, 255), 2)  #可能存在多个人体,提取置信度最高的人体框 def extract_person(image, box_data):     image = np.squeeze(image, axis=0)  # 形状变为 (3, 640, 640)     # 转换为 (H, W, C) 格式     image = np.transpose(image, (1, 2, 0))     print(f"Resized image shape: {image.shape}")      cv2.imshow('read Image',image)     cv2.waitKey(0)     cv2.destroyAllWindows()      boxes = box_data[..., :4].astype(np.int32)     scores = box_data[..., 4]     classes = box_data[..., 5].astype(np.int32)      if len(boxes) > 0:         max_index = np.argmax(scores)         #top, left, right, bottom = boxes[max_index]         left,top,  right, bottom = boxes[max_index]          # 打印原始坐标         print(f"Original Box coordinates: top={top}, left={left}, right={right}, bottom={bottom}")          # 提取目标区域         person = image[top:bottom, left:right]         print(f"Extracted person image size: {person.shape}")          cv2.imshow('extracted person',person)         cv2.waitKey(0)         cv2.destroyAllWindows()          # 保存图像         cv2.imwrite('extract.png', person*255)          return person   if __name__ == "__main__":     onnx_path = r'E:\detect_person\person.onnx'     model = YOLOV5(onnx_path)     output, or_img, resize_img = model.inference(r"G:\depth_detect\huang2\huang1.png")     outbox = filter_box(output, 0.35, 0.35)     if len(outbox) > 0:     	#原图画框可视化         #draw(or_img, outbox)         #提取目标区域         extract = extract_person(resize_img,outbox)     else:         print("No objects detected.")  

相关内容

热门资讯

安卓10系统更新关闭,全面优化... 你知道吗?最近安卓系统又来了一次大动作,那就是安卓10系统的更新关闭了!这可真是让人有点摸不着头脑,...
安卓系统的文件加密,Andro... 你知道吗?在咱们这个数字化时代,保护隐私和安全变得比以往任何时候都重要。尤其是对于安卓系统用户来说,...
使用安卓系统的费用,全面了解使... 你有没有想过,为什么有些人拿着安卓手机,而有些人却选择了苹果?这其中可不仅仅是品牌喜好那么简单,使用...
vivo用原生安卓系统下载,尽... 你有没有发现,现在手机市场真是热闹非凡,各种品牌争奇斗艳,让人眼花缭乱。不过,今天我要给你安利的,可...
安卓系统好用的桌面时钟,实用好... 你有没有发现,手机里的时钟功能有时候比闹钟还重要呢?想象每天早上被它温柔地叫醒,或者在忙碌的工作间隙...
安卓系统导航车载用优盘,安卓车... 你有没有想过,开车的时候,手机导航虽然方便,但有时候屏幕太小,看不清路线?别急,今天就来给你安利一个...
正确使用电池安卓系统,无忧体验 你知道吗?现在这个智能手机时代,电池续航能力可是大家关注的焦点。尤其是安卓系统用户,电池使用得当与否...
玩吧安卓可以和苹果系统,畅享游... 你知道吗?现在这个时代,手机可是我们生活中不可或缺的好伙伴。不管是安卓还是苹果,它们各有各的特色,各...
安卓系统怎么去掉hd,恢复纯净... 你是不是也和我一样,对安卓手机的系统设置充满了好奇?尤其是那个让人眼花缭乱的“HD”标识,有时候看着...
电脑安卓系统性能表,电脑版性能... 你有没有发现,现在手机电脑的操作系统越来越丰富,尤其是安卓系统,简直就像是个万能的小精灵,啥都能干。...
如何玩转机车安卓系统,玩转机车... 你有没有想过,拥有一台酷炫的机车安卓系统,让你的手机瞬间变身成为一辆会跑的摩托车?想象你可以在手机上...
安卓系统网页怎么回顶部,按钮才... 你是不是在使用安卓系统的手机或平板电脑浏览网页时,不小心翻到了页面底部,现在想回到顶部,却有点摸不着...
为什么安卓系统要认证,安卓系统... 你知道吗?安卓系统最近可是掀起了一阵认证热潮,这可不仅仅是简单的更新换代那么简单哦!为什么安卓系统要...
安卓50原生系统手机,功能革新... 你有没有发现,最近你的安卓手机突然变得不一样了?是不是因为它的系统升级到了安卓50原生系统呢?没错,...
安卓永远比不了的系统,永远无法... 你有没有想过,为什么安卓系统永远比不了某些其他系统呢?是不是每次看到那些流畅无阻、功能强大的设备,心...
安卓8怎么升级11系统,解锁新... 你有没有发现,你的安卓手机已经有点儿“老态龙钟”了?别急,别急,今天就来教你怎么给它来个青春焕发的大...
双系统安卓笔记本,开启移动办公... 你有没有想过,一台既能流畅运行安卓应用,又能轻松驾驭Windows系统的笔记本,会是怎样的体验呢?没...
安卓系统调降噪通透软件,打造清... 你有没有发现,最近你的安卓手机在听音乐或者打电话的时候,声音变得超级清晰,仿佛置身于现场?这可不是你...
安卓系统包后缀名,包后缀名背后... 你有没有发现,每次下载安卓应用时,文件名后面总会有那么几个神秘的字母组合,像是“apk”、“jar”...
安卓系统好用的工作软件,盘点十... 你有没有发现,自从你把手机里的安卓系统升级后,工作效率好像提高了不少呢?今天,就让我来给你细细道来,...