Update v5lite.py

Author: ppogg

python_demo/onnxruntime/v5lite.py

@@ -5,26 +5,17 @@
 import onnxruntime as ort
 
 class yolov5_lite():
-    def __init__(self, model_pb_path, label_path, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5):
+    def __init__(self, model_pb_path, label_path, confThreshold=0.5, nmsThreshold=0.5):
         so = ort.SessionOptions()
         so.log_severity_level = 3
         self.net = ort.InferenceSession(model_pb_path, so)
         self.classes = list(map(lambda x: x.strip(), open(label_path, 'r').readlines()))
-        self.num_classes = len(self.classes)
-        anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
-        self.nl = len(anchors)
-        self.na = len(anchors[0]) // 2
-        self.no = self.num_classes + 5
-        self.grid = [np.zeros(1)] * self.nl
-        self.stride = np.array([8., 16., 32.])
-        self.anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(self.nl, -1, 2)
 
         self.confThreshold = confThreshold
         self.nmsThreshold = nmsThreshold
-        self.objThreshold = objThreshold
         self.input_shape = (self.net.get_inputs()[0].shape[2], self.net.get_inputs()[0].shape[3])
 
-    def resize_image(self, srcimg, keep_ratio=True):
+    def letterBox(self, srcimg, keep_ratio=True):
         top, left, newh, neww = 0, 0, self.input_shape[0], self.input_shape[1]
         if keep_ratio and srcimg.shape[0] != srcimg.shape[1]:
             hw_scale = srcimg.shape[0] / srcimg.shape[1]
@@ -43,90 +34,57 @@ def resize_image(self, srcimg, keep_ratio=True):
             img = cv2.resize(srcimg, self.input_shape, interpolation=cv2.INTER_AREA)
         return img, newh, neww, top, left
 
-    def _make_grid(self, nx=20, ny=20):
-        xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
-        return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)
-
     def postprocess(self, frame, outs, pad_hw):
         newh, neww, padh, padw = pad_hw
         frameHeight = frame.shape[0]
         frameWidth = frame.shape[1]
         ratioh, ratiow = frameHeight / newh, frameWidth / neww
-        # Scan through all the bounding boxes output from the network and keep only the
-        # ones with high confidence scores. Assign the box's class label as the class with the highest score.
         classIds = []
         confidences = []
-        box_index = []
         boxes = []
-        outs = outs[outs[:, 4] > self.objThreshold]
         for detection in outs:
-            scores = detection[5:]
-            classId = np.argmax(scores)
-            confidence = scores[classId]
-            if confidence > self.confThreshold:  # and detection[4] > self.objThreshold:
-                center_x = int((detection[0] - padw) * ratiow)
-                center_y = int((detection[1] - padh) * ratioh)
-                width = int(detection[2] * ratiow)
-                height = int(detection[3] * ratioh)
-                left = int(center_x - width / 2)
-                top = int(center_y - height / 2)
+            scores, classId = detection[4], detection[5]
+            if scores > self.confThreshold:  # and detection[4] > self.objThreshold:
+                x1 = int((detection[0] - padw) * ratiow)
+                y1 = int((detection[1] - padh) * ratioh)
+                x2 = int((detection[2] - padw) * ratiow)
+                y2 = int((detection[3] - padh) * ratioh)
                 classIds.append(classId)
-                confidences.append(float(confidence))
-                boxes.append([left, top, width, height])
+                confidences.append(scores)
+                boxes.append([x1, y1, x2, y2])
 
-        # Perform non maximum suppression to eliminate redundant overlapping boxes with
-        # lower confidences.
-        print(boxes)
+        # # Perform non maximum suppression to eliminate redundant overlapping boxes with
+        # # lower confidences.
         indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)
 
-        for i in indices:
-            box_index.append(i[0])
-
-        for i in box_index:
-            box = boxes[i]
-            left = box[0]
-            top = box[1]
-            width = box[2]
-            height = box[3]
-            frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height)
+        for ind in indices:
+            frame = self.drawPred(frame, classIds[ind], confidences[ind], boxes[ind][0], boxes[ind][1], boxes[ind][2], boxes[ind][3])
         return frame
 
-    def drawPred(self, frame, classId, conf, left, top, right, bottom):
+    def drawPred(self, frame, classId, conf, x1, y1, x2, y2):
         # Draw a bounding box.
-        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=2)
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), thickness=2)
 
         label = '%.2f' % conf
-        label = '%s:%s' % (self.classes[classId], label)
+        text = '%s:%s' % (self.classes[int(classId)], label)
 
         # Display the label at the top of the bounding box
-        labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        top = max(top, labelSize[1])
-        # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED)
-        cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 255, 0), thickness=1)
+        labelSize, baseLine = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        y1 = max(y1, labelSize[1])
+        cv2.putText(frame, text, (x1, y1 - 10), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 255, 0), thickness=1)
         return frame
 
     def detect(self, srcimg):
-        img, newh, neww, top, left = self.resize_image(srcimg)
+        img, newh, neww, top, left = self.letterBox(srcimg)
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         img = img.astype(np.float32) / 255.0
         blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
 
         t1 = time.time()
-        outs = self.net.run(None, {self.net.get_inputs()[0].name: blob})[0].squeeze(axis=0)
+        outs = self.net.run(None, {self.net.get_inputs()[0].name: blob})[0]
         cost_time = time.time() - t1
         print(outs.shape)
-        row_ind = 0
-        for i in range(self.nl):
-            h, w = int(self.input_shape[0] / self.stride[i]), int(self.input_shape[1] / self.stride[i])
-            length = int(self.na * h * w)
-            if self.grid[i].shape[2:4] != (h, w):
-                self.grid[i] = self._make_grid(w, h)
-
-            outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
-                self.grid[i], (self.na, 1))) * int(self.stride[i])
-            outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat(
-                self.anchor_grid[i], h * w, axis=0)
-            row_ind += length
+
         srcimg = self.postprocess(srcimg, outs, (newh, neww, top, left))
         infer_time = 'Inference Time: ' + str(int(cost_time * 1000)) + 'ms'
         cv2.putText(srcimg, infer_time, (5, 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 0, 0), thickness=1)
@@ -135,8 +93,8 @@ def detect(self, srcimg):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--imgpath', type=str, default='../sample/horse.jpg', help="image path")
-    parser.add_argument('--modelpath', type=str, default='../weights/v5Lite-e.onnx', help="onnx filepath")
+    parser.add_argument('--imgpath', type=str, default='./000000001000.jpg', help="image path")
+    parser.add_argument('--modelpath', type=str, default='./v5lite-e_end2end.onnx', help="onnx filepath")
     parser.add_argument('--classfile', type=str, default='coco.names', help="classname filepath")
     parser.add_argument('--confThreshold', default=0.5, type=float, help='class confidence')
     parser.add_argument('--nmsThreshold', default=0.6, type=float, help='nms iou thresh')
@@ -146,9 +104,4 @@ def detect(self, srcimg):
     net = yolov5_lite(args.modelpath, args.classfile, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold)
     srcimg = net.detect(srcimg.copy())
 
-    winName = 'Deep learning object detection in onnxruntime'
-    cv2.namedWindow(winName, cv2.WINDOW_NORMAL)
-    cv2.imshow(winName, srcimg)
-    cv2.waitKey(0)
-    # cv2.imwrite('save.jpg', srcimg )
-    cv2.destroyAllWindows()
+    cv2.imwrite('save.jpg', srcimg )