File size: 3,254 Bytes
963830b a07f7bd d2b2b68 963830b e363419 963830b e363419 963830b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import numpy as np
import cv2
class YoloX:
def __init__(self, modelPath, confThreshold=0.35, nmsThreshold=0.5, objThreshold=0.5, backendId=0, targetId=0):
self.num_classes = 80
self.net = cv2.dnn.readNet(modelPath)
self.input_size = (640, 640)
self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
self.strides = [8, 16, 32]
self.confThreshold = confThreshold
self.nmsThreshold = nmsThreshold
self.objThreshold = objThreshold
self.backendId = backendId
self.targetId = targetId
self.net.setPreferableBackend(self.backendId)
self.net.setPreferableTarget(self.targetId)
self.generateAnchors()
@property
def name(self):
return self.__class__.__name__
def setBackendAndTarget(self, backendId, targetId):
self.backendId = backendId
self.targetId = targetId
self.net.setPreferableBackend(self.backendId)
self.net.setPreferableTarget(self.targetId)
def preprocess(self, img):
blob = np.transpose(img, (2, 0, 1))
return blob[np.newaxis, :, :, :]
def infer(self, srcimg):
input_blob = self.preprocess(srcimg)
self.net.setInput(input_blob)
outs = self.net.forward(self.net.getUnconnectedOutLayersNames())
predictions = self.postprocess(outs[0])
return predictions
def postprocess(self, outputs):
dets = outputs[0]
dets[:, :2] = (dets[:, :2] + self.grids) * self.expanded_strides
dets[:, 2:4] = np.exp(dets[:, 2:4]) * self.expanded_strides
# get boxes
boxes = dets[:, :4]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
# get scores and class indices
scores = dets[:, 4:5] * dets[:, 5:]
max_scores = np.amax(scores, axis=1)
max_scores_idx = np.argmax(scores, axis=1)
keep = cv2.dnn.NMSBoxesBatched(boxes_xyxy.tolist(), max_scores.tolist(), max_scores_idx.tolist(), self.confThreshold, self.nmsThreshold)
candidates = np.concatenate([boxes_xyxy, max_scores[:, None], max_scores_idx[:, None]], axis=1)
if len(keep) == 0:
return np.array([])
return candidates[keep]
def generateAnchors(self):
self.grids = []
self.expanded_strides = []
hsizes = [self.input_size[0] // stride for stride in self.strides]
wsizes = [self.input_size[1] // stride for stride in self.strides]
for hsize, wsize, stride in zip(hsizes, wsizes, self.strides):
xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))
grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
self.grids.append(grid)
shape = grid.shape[:2]
self.expanded_strides.append(np.full((*shape, 1), stride))
self.grids = np.concatenate(self.grids, 1)
self.expanded_strides = np.concatenate(self.expanded_strides, 1)
|