Update 部署 OpenVINO YOLOv5推理

Author: JasonSloan

03. 部署/02. OpenVINO/02. YOLOv5独立运行的推理代码/standalone_inference.py

@@ -0,0 +1,82 @@
+import os
+import sys
+from pathlib import Path
+
+import torch
+import cv2
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+# from models.common import DetectMultiBackend
+from standalone_model import DetectMultiBackend
+from standalone_utils import increment_path, colorstr, Profile, LOGGER, LoadImages, scale_boxes, non_max_suppression
+
+
+@torch.no_grad()
+def run(
+        weights=ROOT / 'yolov5s/POT_INT8_openvino_model',  # model path or triton URL
+        source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
+        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+        imgsz=(640, 640),  # inference size (height, width)
+        conf_thres=0.25,  # confidence threshold
+        iou_thres=0.45,  # NMS IOU threshold
+        max_det=1000,  # maximum detections per image
+        device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+        save_img=True,  # save results
+        classes=None,  # filter by class: --class 0, or --class 0 2 3
+        agnostic_nms=False,  # class-agnostic NMS
+        augment=False,  # augmented inference
+        visualize=False,  # visualize features
+        project=ROOT / 'runs/detect',  # save results to project/name
+        name='exp',  # save results to project/name
+        exist_ok=False,  # existing project/name ok, do not increment
+        line_thickness=3,  # bounding box thickness (pixels)
+        hide_labels=False,  # hide labels
+        hide_conf=False,  # hide confidences
+        half=False,  # use FP16 half-precision inference
+        vid_stride=1,  # video frame-rate stride
+):
+    source = str(source)
+
+    # Directories
+    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
+    save_dir.mkdir(parents=True, exist_ok=True)  # make dir
+
+    # Load model
+    device = torch.device(device)
+    model = DetectMultiBackend(weights, device=device, data=data, fp16=half)
+    stride, names = model.stride, model.names
+
+    # Dataloader
+    bs = 1  # batch_size
+    # 前处理只有letterbox+HWC2CHW+BGR2RGB+ascontiguousarray
+    dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=False, vid_stride=vid_stride)
+    vid_path, vid_writer = [None] * bs, [None] * bs
+
+    # Run inference
+    seen, dt = 0, (Profile(), Profile(), Profile())
+    for path, im, im0s, vid_cap, s in dataset:
+        with dt[0]:
+            im = torch.from_numpy(im).to(model.device)
+            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+            im /= 255  # 0 - 255 to 0.0 - 1.0
+            im = im[None]  # expand for batch dim
+
+        # Inference
+        with dt[1]:
+            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
+            pred = model(im, augment=augment, visualize=visualize)
+
+        # NMS
+        with dt[2]:
+            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
+        return pred    
+  
+
+
+if __name__ == "__main__":
+    run()

03. 部署/02. OpenVINO/02. YOLOv5独立运行的推理代码/standalone_model.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+import yaml
+import torch
+import torch.nn as nn
+import numpy as np
+from openvino.runtime import Core, Layout, get_batch
+from standalone_utils import LOGGER
+
+
+class DetectMultiBackend(nn.Module):
+    # YOLOv5 MultiBackend class for python inference on various backends
+    def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
+        super().__init__()
+        w = str(weights[0] if isinstance(weights, list) else weights)
+        LOGGER.info(f'Loading {w} for OpenVINO inference...')
+        ie = Core()
+        w = next(Path(w).glob('*.xml'))  # get *.xml file from *_openvino_model dir
+        network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin'))
+        network.get_parameters()[0].set_layout(Layout("NCHW"))
+        batch_dim = get_batch(network)
+        batch_size = batch_dim.get_length()
+        executable_network = ie.compile_model(network, device_name="CPU")  # device_name="MYRIAD" for Intel NCS2
+        stride, names = self._load_metadata(Path(w).with_suffix('.yaml'))  # load metadata
+        self.__dict__.update(locals())  # assign all variables to self
+
+    def forward(self, im, augment=False, visualize=False):
+        im = im.cpu().numpy()  # FP32
+        y = list(self.executable_network([im]).values())
+        return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
+    
+    def from_numpy(self, x):
+        return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
+        
+    def _load_metadata(self, file=Path('path/to/meta.yaml')):
+        with open(file, errors='ignore') as f:
+            d = yaml.safe_load(f)
+        return d['stride'], d['names']  # assign stride, names