import torch import cv2 import numpy as np import torch import os import importlib from model.plugins.ModelBase import ModelBase from loguru import logger ''' class ModelManager_tmp(): def __init__(self): print("ModelInit") def __del__(self): print("ModelManager DEL") def __preprocess_image(self,image, cfg, bgr2rgb=True): """图片预处理""" img, scale_ratio, pad_size = letterbox(image, new_shape=cfg['input_shape']) if bgr2rgb: img = img[:, :, ::-1] img = img.transpose(2, 0, 1) # HWC2CHW img = np.ascontiguousarray(img, dtype=np.float32) return img, scale_ratio, pad_size def __draw_bbox(self,bbox, img0, color, wt, names): """在图片上画预测框""" det_result_str = '' for idx, class_id in enumerate(bbox[:, 5]): if float(bbox[idx][4] < float(0.05)): continue img0 = cv2.rectangle(img0, (int(bbox[idx][0]), int(bbox[idx][1])), (int(bbox[idx][2]), int(bbox[idx][3])), color, wt) img0 = cv2.putText(img0, str(idx) + ' ' + names[int(class_id)], (int(bbox[idx][0]), int(bbox[idx][1] + 16)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) img0 = cv2.putText(img0, '{:.4f}'.format(bbox[idx][4]), (int(bbox[idx][0]), int(bbox[idx][1] + 32)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) det_result_str += '{} {} {} {} {} {}\n'.format( names[bbox[idx][5]], str(bbox[idx][4]), bbox[idx][0], bbox[idx][1], bbox[idx][2], bbox[idx][3]) return img0 def __get_labels_from_txt(self,path): """从txt文件获取图片标签""" labels_dict = dict() with open(path) as f: for cat_id, label in enumerate(f.readlines()): labels_dict[cat_id] = label.strip() return labels_dict def __draw_prediction(self,pred, image, labels): """在图片上画出预测框并进行可视化展示""" imgbox = widgets.Image(format='jpg', height=720, width=1280) img_dw = self.__draw_bbox(pred, image, (0, 255, 0), 2, labels) imgbox.value = cv2.imencode('.jpg', img_dw)[1].tobytes() display(imgbox) def __infer_image(self,img_path, model, class_names, cfg): """图片推理""" # 图片载入 image = cv2.imread(img_path) # 数据预处理 img, scale_ratio, pad_size = self.__preprocess_image(image, cfg) # 模型推理 output = model.infer([img])[0] output = torch.tensor(output) # 非极大值抑制后处理 boxout = nms(output, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"]) pred_all = boxout[0].numpy() # 预测坐标转换 scale_coords(cfg['input_shape'], pred_all[:, :4], image.shape, ratio_pad=(scale_ratio, pad_size)) # 图片预测结果可视化 self.__draw_prediction(pred_all, image, class_names) def __infer_frame_with_vis(self,image, model, labels_dict, cfg, bgr2rgb=True): # 数据预处理 img, scale_ratio, pad_size = self.__preprocess_image(image, cfg, bgr2rgb) # 模型推理 output = model.infer([img])[0] output = torch.tensor(output) # 非极大值抑制后处理 boxout = nms(output, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"]) pred_all = boxout[0].numpy() # 预测坐标转换 scale_coords(cfg['input_shape'], pred_all[:, :4], image.shape, ratio_pad=(scale_ratio, pad_size)) # 图片预测结果可视化 img_vis = self.__draw_bbox(pred_all, image, (0, 255, 0), 2, labels_dict) return img_vis def __img2bytes(self,image): """将图片转换为字节码""" return bytes(cv2.imencode('.jpg', image)[1]) def __infer_camera(self,model, labels_dict, cfg): """外设摄像头实时推理""" def find_camera_index(): max_index_to_check = 10 # Maximum index to check for camera for index in range(max_index_to_check): cap = cv2.VideoCapture(index) if cap.read()[0]: cap.release() return index # If no camera is found raise ValueError("No camera found.") # 获取摄像头 --这里可以换成RTSP流 camera_index = find_camera_index() cap = cv2.VideoCapture(camera_index) # 初始化可视化对象 image_widget = widgets.Image(format='jpeg', width=1280, height=720) display(image_widget) while True: # 对摄像头每一帧进行推理和可视化 _, img_frame = cap.read() image_pred = self.__infer_frame_with_vis(img_frame, model, labels_dict, cfg) image_widget.value = self.__img2bytes(image_pred) def __infer_video(self,video_path, model, labels_dict, cfg): """视频推理""" image_widget = widgets.Image(format='jpeg', width=800, height=600) display(image_widget) # 读入视频 cap = cv2.VideoCapture(video_path) while True: ret, img_frame = cap.read() if not ret: break # 对视频帧进行推理 image_pred = self.__infer_frame_with_vis(img_frame, model, labels_dict, cfg, bgr2rgb=True) image_widget.value = self.__img2bytes(image_pred) def startWork(self,infer_mode,file_paht = ""): cfg = { 'conf_thres': 0.4, # 模型置信度阈值，阈值越低，得到的预测框越多 'iou_thres': 0.5, # IOU阈值，高于这个阈值的重叠预测框会被过滤掉 'input_shape': [640, 640], # 模型输入尺寸 } model_path = 'yolo.om' label_path = './coco_names.txt' # 初始化推理模型 model = InferSession(0, model_path) labels_dict = self.__get_labels_from_txt(label_path) #执行验证 if infer_mode == 'image': img_path = 'world_cup.jpg' self.__infer_image(img_path, model, labels_dict, cfg) elif infer_mode == 'camera': self.__infer_camera(model, labels_dict, cfg) elif infer_mode == 'video': video_path = 'racing.mp4' self.__infer_video(video_path, model, labels_dict, cfg) ''' ''' 算法实现类，实现算法执行线程，根据配内容，以线程方式执行算法模块 ''' class ModelManager(): def __init__(self): print("ModelManager init") def __del__(self): print("ModelManager del") def doWork(self): pass #动态导入文件 -- 方法二 -- 相对推荐使用该方法 但spec感觉没什么用 def import_source(spec, plgpath): module = None if os.path.exists(plgpath): module_spec = importlib.util.spec_from_file_location(spec, plgpath) module = importlib.util.module_from_spec(module_spec) module_spec.loader.exec_module(module) else: logger.error("{}文件不存在".format(plgpath)) return module #plgpath 为list [poc][file_name][name] def run_plugin(plgpath, target,copy_flag=True): module = import_source("", plgpath) if module: classname = "Model" plg = getattr(module, classname)() if not isinstance(plg, ModelBase): raise Exception("{} not rx_Model".format(plg)) new_plg = plg result = new_plg.doWork("","","","") # 执行plugin基类的run, 返回结果 return result else: print("模型加载失败") return None def test(): run_plugin("plugins/RYRQ_Model_ACL.py","") if __name__ == "__main__": test()