| import numpy as np | |
| import cv2 | |
| class yolox(): | |
| def __init__(self, model, p6=False, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5): | |
| with open('coco.names', 'rt') as f: | |
| self.class_names = f.read().rstrip('\n').split('\n') | |
| self.net = cv2.dnn.readNet(model) | |
| self.input_size = (640, 640) | |
| self.mean = (0.485, 0.456, 0.406) | |
| self.std = (0.229, 0.224, 0.225) | |
| if not p6: | |
| self.strides = [8, 16, 32] | |
| else: | |
| self.strides = [8, 16, 32, 64] | |
| self.confThreshold = confThreshold | |
| self.nmsThreshold = nmsThreshold | |
| self.objThreshold = objThreshold | |
| def preprocess(self, image): | |
| if len(image.shape) == 3: | |
| padded_img = np.ones((self.input_size[0], self.input_size[1], 3)) * 114.0 | |
| else: | |
| padded_img = np.ones(self.input_size) * 114.0 | |
| img = np.array(image) | |
| r = min(self.input_size[0] / img.shape[0], self.input_size[1] / img.shape[1]) | |
| resized_img = cv2.resize( | |
| img, (int(img.shape[1] * r), int(img.shape[0] * r)), interpolation=cv2.INTER_LINEAR | |
| ).astype(np.float32) | |
| padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img | |
| image = padded_img | |
| image = image.astype(np.float32) | |
| image = image[:, :, ::-1] | |
| image /= 255.0 | |
| image -= self.mean | |
| image /= self.std | |
| return image, r | |
| def demo_postprocess(self, outputs): | |
| grids = [] | |
| 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) | |
| grids.append(grid) | |
| shape = grid.shape[:2] | |
| expanded_strides.append(np.full((*shape, 1), stride)) | |
| grids = np.concatenate(grids, 1) | |
| expanded_strides = np.concatenate(expanded_strides, 1) | |
| outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides | |
| outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides | |
| return outputs | |
| def nms(self, boxes, scores): | |
| """Single class NMS implemented in Numpy.""" | |
| x1 = boxes[:, 0] | |
| y1 = boxes[:, 1] | |
| x2 = boxes[:, 2] | |
| y2 = boxes[:, 3] | |
| areas = (x2 - x1 + 1) * (y2 - y1 + 1) | |
| order = scores.argsort()[::-1] | |
| keep = [] | |
| while order.size > 0: | |
| i = order[0] | |
| keep.append(i) | |
| xx1 = np.maximum(x1[i], x1[order[1:]]) | |
| yy1 = np.maximum(y1[i], y1[order[1:]]) | |
| xx2 = np.minimum(x2[i], x2[order[1:]]) | |
| yy2 = np.minimum(y2[i], y2[order[1:]]) | |
| w = np.maximum(0.0, xx2 - xx1 + 1) | |
| h = np.maximum(0.0, yy2 - yy1 + 1) | |
| inter = w * h | |
| ovr = inter / (areas[i] + areas[order[1:]] - inter) | |
| inds = np.where(ovr <= self.nmsThreshold)[0] | |
| order = order[inds + 1] | |
| return keep | |
| def multiclass_nms(self, boxes, scores): | |
| """Multiclass NMS implemented in Numpy""" | |
| final_dets = [] | |
| num_classes = scores.shape[1] | |
| for cls_ind in range(num_classes): | |
| cls_scores = scores[:, cls_ind] | |
| valid_score_mask = cls_scores > self.confThreshold | |
| if valid_score_mask.sum() == 0: | |
| continue | |
| else: | |
| valid_scores = cls_scores[valid_score_mask] | |
| valid_boxes = boxes[valid_score_mask] | |
| keep = self.nms(valid_boxes, valid_scores) | |
| if len(keep) > 0: | |
| cls_inds = np.ones((len(keep), 1)) * cls_ind | |
| dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1) | |
| final_dets.append(dets) | |
| if len(final_dets) == 0: | |
| return None | |
| return np.concatenate(final_dets, 0) | |
| def vis(self, img, boxes, scores, cls_ids): | |
| detected_classes = [] | |
| for i in range(len(boxes)): | |
| box = boxes[i] | |
| cls_id = int(cls_ids[i]) | |
| score = scores[i] | |
| if score < self.confThreshold: | |
| continue | |
| x0 = int(box[0]) | |
| y0 = int(box[1]) | |
| x1 = int(box[2]) | |
| y1 = int(box[3]) | |
| class_name = self.class_names[cls_id] | |
| detected_classes.append(class_name) | |
| text = '{}:{:.1f}%'.format(class_name, score * 100) | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| txt_size = cv2.getTextSize(text, font, 0.4, 1)[0] | |
| cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), 2) | |
| cv2.rectangle(img, (x0, y0 + 1), (x0 + txt_size[0] + 1, y0 + int(1.5 * txt_size[1])), (255, 255, 255), -1) | |
| cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, (0, 0, 0), thickness=1) | |
| return img, detected_classes | |
| def detect(self, srcimg): | |
| detected_classes = [] | |
| img, ratio = self.preprocess(srcimg) | |
| blob = cv2.dnn.blobFromImage(img) | |
| self.net.setInput(blob) | |
| outs = self.net.forward(self.net.getUnconnectedOutLayersNames()) | |
| predictions = self.demo_postprocess(outs[0])[0] | |
| boxes = predictions[:, :4] | |
| scores = predictions[:, 4:5] * predictions[:, 5:] | |
| 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. | |
| boxes_xyxy /= ratio | |
| dets = self.multiclass_nms(boxes_xyxy, scores) | |
| if dets is not None: | |
| final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] | |
| srcimg, detected_classes = self.vis(srcimg, final_boxes, final_scores, final_cls_inds) | |
| return srcimg, ", ".join(list(set(detected_classes))) if len(detected_classes) > 0 else "" | |