import cv2 import numpy as np import axengine as axe import matplotlib import argparse class Colors: def __init__(self): self.palette = [self.hex2rgb(c) for c in matplotlib.colors.TABLEAU_COLORS.values()] self.n = len(self.palette) def __call__(self, i, bgr=False): c = self.palette[int(i) % self.n] return (c[2], c[1], c[0]) if bgr else c @staticmethod def hex2rgb(h): return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) colors = Colors() def plot_one_box(x, im, color=None, label=None, line_thickness=3, kpt_label=False, kpts=None, steps=2, orig_shape=None): assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.' tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(im, c1, c2, color, thickness=tl*1//3, lineType=cv2.LINE_AA) if label: if len(label.split(' ')) > 1: tf = max(tl - 1, 1) t_size = cv2.getTextSize(label, 0, fontScale=tl / 6, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 6, [225, 255, 255], thickness=tf//2, lineType=cv2.LINE_AA) if kpt_label: plot_skeleton_kpts(im, kpts, steps, orig_shape=orig_shape) def plot_skeleton_kpts(im, kpts, steps, orig_shape=None): palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255], [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255], [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102], [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]]) num_kpts = len(kpts) // steps skeleton = [[1, 2], [2, 3], [14, 1], [14, 4], [4, 5], [5, 6], [13, 14], [7, 14], [10, 14], [7, 8], [8, 9],[10,11],[11, 12]] pose_limb_color = palette[[9, 9, 9, 9, 9, 9, 7, 0, 0, 0, 0, 0, 0]] pose_kpt_color = palette[[16, 16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9]] radius = 5 for kid in range(num_kpts): r, g, b = pose_kpt_color[kid] x_coord, y_coord = kpts[steps * kid], kpts[steps * kid + 1] if not (x_coord % 640 == 0 or y_coord % 640 == 0): if steps == 3: conf = kpts[steps * kid + 2] cv2.circle(im, (int(x_coord), int(y_coord)), radius, (int(r), int(g), int(b)), -1) cv2.putText(im,str(kid),(int(x_coord-2), int(y_coord-2)),cv2.FONT_HERSHEY_COMPLEX_SMALL, 1,(0,0,255),1) for sk_id, sk in enumerate(skeleton): r, g, b = pose_limb_color[sk_id] pos1 = (int(kpts[(sk[0]-1)*steps]), int(kpts[(sk[0]-1)*steps+1])) pos2 = (int(kpts[(sk[1]-1)*steps]), int(kpts[(sk[1]-1)*steps+1])) if steps == 3: conf1 = kpts[(sk[0]-1)*steps+2] conf2 = kpts[(sk[1]-1)*steps+2] if pos1[0]%640 == 0 or pos1[1]%640==0 or pos1[0]<0 or pos1[1]<0: continue if pos2[0] % 640 == 0 or pos2[1] % 640 == 0 or pos2[0]<0 or pos2[1]<0: continue cv2.line(im, pos1, pos2, (int(r), int(g), int(b)), thickness=2) def box_iou(box1, box2, eps=1e-7): (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2) inter = (np.min(a2, b2) - np.max(a1, b1)).clamp(0).prod(2) return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps) def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32): shape = im.shape[:2] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: r = min(r, 1.0) ratio = r, r new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] if auto: dw, dh = np.mod(dw, stride), np.mod(dh, stride) elif scaleFill: dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] dw /= 2 dh /= 2 if shape[::-1] != new_unpad: im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) return im, ratio, (dw, dh) def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None, kpt_label=False, step=2): if ratio_pad is None: gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 else: gain = ratio_pad[0] pad = ratio_pad[1] if isinstance(gain, (list, tuple)): gain = gain[0] if not kpt_label: coords[:, [0, 2]] -= pad[0] coords[:, [1, 3]] -= pad[1] coords[:, [0, 2]] /= gain coords[:, [1, 3]] /= gain else: coords[:, 0::step] -= pad[0] coords[:, 1::step] -= pad[1] coords[:, 0::step] /= gain coords[:, 1::step] /= gain return coords def clip_coords(boxes, img_shape, step=2): boxes[:, 0::step].clamp_(0, img_shape[1]) boxes[:, 1::step].clamp_(0, img_shape[0]) def model_inference(model_path=None, input=None): session = axe.InferenceSession(model_path, None) input_name = session.get_inputs()[0].name output = session.run(None, {input_name: input}) return output def xywh2xyxy(x): y = np.copy(x) y[..., 0] = x[..., 0] - x[..., 2] / 2 y[..., 1] = x[..., 1] - x[..., 3] / 2 y[..., 2] = x[..., 0] + x[..., 2] / 2 y[..., 3] = x[..., 1] + x[..., 3] / 2 return y def nms_boxes(boxes, scores): x = boxes[:, 0] y = boxes[:, 1] w = boxes[:, 2] - boxes[:, 0] h = boxes[:, 3] - boxes[:, 1] areas = w * h order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) xx1 = np.maximum(x[i], x[order[1:]]) yy1 = np.maximum(y[i], y[order[1:]]) xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) inter = w1 * h1 ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= 0.45)[0] order = order[inds + 1] keep = np.array(keep) return keep def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, labels=(), kpt_label=False, nc=None, nkpt=14): if nc is None: nc = prediction.shape[2] - 5 if not kpt_label else prediction.shape[2] - (5+3*nkpt) xc = prediction[..., 4] > conf_thres min_wh, max_wh = 2, 4096 max_det = 300 max_nms = 30000 redundant = True multi_label &= nc > 1 merge = False output = [np.zeros((0,6))] * prediction.shape[0] for xi, x in enumerate(prediction): x = x[xc[xi]] if labels and len(labels[xi]): l = labels[xi] v = np.zeros(len(l), nc + 5) v[:, :4] = l[:, 1:5] v[:, 4] = 1.0 v[range(len(l)), l[:, 0].long() + 5] = 1.0 x = np.concatenate((x, v), 0) if not x.shape[0]: continue x[:, 5:5+nc] *= x[:, 4:5] box = xywh2xyxy(x[:, :4]) if multi_label: if not kpt_label: i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T x = np.concatenate((box[i], x[i, j + 5, None], j[:, None].float()), 1) else: kpts = x[:, 5+nc:] i, j = (x[:, 5:5+nc] > conf_thres).nonzero(as_tuple=False).T x = np.concatenate((box[i], x[i, j + 5, None], j[:, None].float(),kpts[i]), 1) else: if not kpt_label: conf, j = x[:, 5:].max(1, keepdim=True) x = np.concatenate((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] else: kpts = x[:, 5+nc:] conf = np.max(x[:, 5:5+nc], 1).reshape(box.shape[:1][0], 1) j = np.argmax(x[:, 5:5+nc], 1).reshape(box.shape[:1][0], 1) x = np.concatenate((box, conf, j, kpts), 1)[conf.reshape(box.shape[:1][0]) > conf_thres] if classes is not None: x = x[(x[:, 5:6] == np.array(classes, device=x.device)).any(1)] n = x.shape[0] if not n: continue elif n > max_nms: x = x[x[:, 4].argsort(descending=True)[:max_nms]] c = x[:, 5:6] * (0 if agnostic else max_wh) boxes, scores = x[:, :4] + c, x[:, 4] i = nms_boxes(boxes, scores) if i.shape[0] > max_det: i = i[:max_det] if merge and (1 < n < 3E3): iou = box_iou(boxes[i], boxes) > iou_thres weights = iou * scores[None] x[i, :4] = np.multiply(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) if redundant: i = i[iou.sum(1) > 1] output[xi] = x[i] return output def _make_grid(nx=20, ny=20): y, x = np.arange(ny, dtype=np.float32), np.arange(nx, dtype=np.float32) yv, xv = np.meshgrid(y, x, indexing='ij') return np.stack((xv, yv), 2).reshape((1, 1, ny, nx, 2)) def sigmoid(x): return 1 / (1 + np.exp(-x)) def preprocess(img_path, imgsz): """预处理：读取图像并进行归一化""" im0 = cv2.imread(img_path) img = letterbox(im0, imgsz, auto=False, stride=32)[0] img = np.ascontiguousarray(img[:, :, ::-1].transpose(2, 0, 1)) img = np.asarray(img, dtype=np.uint8) img = np.expand_dims(img, 0) return img, im0 def model_postprocess(preds, anchors, stride, names, nkpt, conf_thres, iou_thres): """后处理：解码预测结果、NMS和坐标变换""" na = len(anchors[0]) // 2 nl = len(anchors) nc = len(names) no = len(names) + 5 + nkpt * 3 z = [] for i, pred in enumerate(preds): bs, _, ny, nx = pred.shape pred = pred.reshape(bs, na, no, ny, nx).transpose(0, 1, 3, 4, 2) pred_det = pred[..., :5+nc] pred_kpt = pred[..., 5+nc:] grid = _make_grid(nx, ny) kpt_grid_x = grid[..., 0:1] kpt_grid_y = grid[..., 1:2] y = sigmoid(pred_det) xy = (y[..., 0:2] * 2. - 0.5 + grid) * stride[i] wh = (y[..., 2:4] * 2) ** 2 * np.array(anchors[i]).reshape(1, 3, 1, 1, 2) pred_kpt[..., 0::3] = (pred_kpt[..., ::3] * 2. - 0.5 + np.tile(kpt_grid_x, (1,1,1,1,nkpt))) * stride[i] pred_kpt[..., 1::3] = (pred_kpt[..., 1::3] * 2. - 0.5 + np.tile(kpt_grid_y, (1,1,1,1,nkpt))) * stride[i] pred_kpt[..., 2::3] = sigmoid(pred_kpt[..., 2::3]) y = np.concatenate((xy, wh, y[..., 4:], pred_kpt), axis=-1) z.append(y.reshape(bs, na * nx * ny, no)) preds = np.concatenate(z, 1) preds = non_max_suppression(preds, conf_thres, iou_thres, nc=nc, nkpt=nkpt, kpt_label=True) return preds def draw_predictions(preds, img, im0, names, imgsz): """绘制检测结果和关键点""" for i, det in enumerate(preds): if len(det): scale_coords(imgsz, det[:, :4], im0.shape, kpt_label=False) scale_coords(imgsz, det[:, 6:], im0.shape, kpt_label=True, step=3) for det_index, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])): print("class:",names[int(cls)], "left:%.0f" % xyxy[0],"top:%.0f" % xyxy[1],"right:%.0f" % xyxy[2],"bottom:%.0f" % xyxy[3], "conf:",'{:.0f}%'.format(float(conf)*100)) c = int(cls) label = f'{names[c]} {conf:.2f}' kpts = det[det_index, 6:] plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=2, kpt_label=True, kpts=kpts, steps=3, orig_shape=im0.shape[:2]) return im0 if __name__ == "__main__": parser = argparse.ArgumentParser(description='跌倒检测模型推理脚本') parser.add_argument('--model', type=str, default='./fall_ax650_npu3.axmodel', help='axmodel 模型路径') parser.add_argument('--img', type=str, default='./fall4.png', help='输入图像路径') parser.add_argument('--output', type=str, default='axmodel_res.jpg', help='输出结果图像路径') parser.add_argument('--imgsz', type=int, nargs=2, default=[320, 480], help='输入图像尺寸 (height width)') parser.add_argument('--conf-thres', type=float, default=0.3, help='置信度阈值') parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU阈值') args = parser.parse_args() # model params names = ['normal', 'fall'] anchors = [[30, 61, 55, 124, 90, 207], [149, 232, 128, 357, 221, 308]] stride = [16, 32] nkpt = 14 imgsz = tuple(args.imgsz) img, im0 = preprocess(args.img, imgsz) preds = model_inference(args.model, img) preds = model_postprocess(preds, anchors, stride, names, nkpt, args.conf_thres, args.iou_thres) im0 = draw_predictions(preds, img, im0, names, imgsz) cv2.imwrite(args.output, im0) print(f"Result saved to {args.output}")