Spaces:
Running
Running
| 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() | |
| 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) | |