python版本基于瑞芯微rk3588 NPU的NanoTrack跟踪算法源码+项目操作说明(运行可达120FPS).zip

from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals import numpy as np import cv2 from core.config import cfg from utils.bbox import corner2center from rknnlite.api import RKNNLite class NnoTracker_RKNNLite(object): def __init__(self,Tback_weight, Xback_weight, Head_weight): self.score_size = (cfg.TRACK.INSTANCE_SIZE - cfg.TRACK.EXEMPLAR_SIZE) // \ cfg.POINT.STRIDE + 1 + cfg.TRACK.BASE_SIZE hanning = np.hanning(self.score_size) window = np.outer(hanning, hanning) self.cls_out_channels = 2 self.window = window.flatten() self.points = self.generate_points(cfg.POINT.STRIDE, self.score_size) #--------------------------------------------------------# #--------------modify environment------------------------# # 1. T init self.rknn_Tback = RKNNLite() # load RKNN model print('--> Load RKNN model') ret = self.rknn_Tback.load_rknn(Tback_weight) if ret != 0: print('Load RKNN model failed') exit(ret) print('done') # init runtime environment print('--> Init runtime environment') ret = self.rknn_Tback.init_runtime(core_mask=RKNNLite.NPU_CORE_0) if ret != 0: print('Init runtime environment failed') exit(ret) print('done') # 2. X init self.rknn_Xback = RKNNLite() # Load model print('--> rknn_Xback: Loading model') ret = self.rknn_Xback.load_rknn(Xback_weight) if ret != 0: print('rknn_Xback: Load model failed!') exit(ret) print('rknn_Xback:done') # Init runtime environment print('--> Init runtime environment') ret = self.rknn_Xback.init_runtime(core_mask=RKNNLite.NPU_CORE_1) if ret != 0: print('Init runtime environment failed!') exit(ret) print('done') # 3. Head init self.rknn_Head = RKNNLite() # Load model print('--> rknn_Head: Loading model') ret = self.rknn_Head.load_rknn(Head_weight) if ret != 0: print('rknn_Head: Load model failed!') exit(ret) print('rknn_Head:done') # Init runtime environment print('--> Init runtime environment') ret = self.rknn_Head.init_runtime(core_mask=RKNNLite.NPU_CORE_2) if ret != 0: print('Init runtime environment failed!') exit(ret) print('done') def generate_points(self, stride, size): ori = - (size // 2) * stride x, y = np.meshgrid([ori + stride * dx for dx in np.arange(0, size)], [ori + stride * dy for dy in np.arange(0, size)]) points = np.zeros((size * size, 2), dtype=np.float32) points[:, 0], points[:, 1] = x.astype(np.float32).flatten(), y.astype(np.float32).flatten() return points def _convert_bbox(self, delta, point): delta = delta.permute(1, 2, 3, 0).contiguous().view(4, -1) delta = delta.detach().cpu().numpy() delta[0, :] = point[:, 0] - delta[0, :] # x1 delta[1, :] = point[:, 1] - delta[1, :] # y1 delta[2, :] = point[:, 0] + delta[2, :] # x2 delta[3, :] = point[:, 1] + delta[3, :] # y2 delta[0, :], delta[1, :], delta[2, :], delta[3, :] = corner2center(delta) return delta def _convert_score(self, score): if self.cls_out_channels == 1: score = score.permute(1, 2, 3, 0).contiguous().view(-1) score = score.sigmoid().detach().cpu().numpy() else: score = score.permute(1, 2, 3, 0).contiguous().view(self.cls_out_channels, -1).permute(1, 0) score = score.softmax(1).detach()[:, 1].cpu().numpy() return score def _convert_bbox_numpy(self, delta, point): # delta = delta.permute(1, 2, 3, 0).contiguous().view(4, -1) # delta = delta.detach().cpu().numpy() delta = delta.transpose((1,2,3,0)).reshape(4, -1) delta[0, :] = point[:, 0] - delta[0, :] # x1 delta[1, :] = point[:, 1] - delta[1, :] # y1 delta[2, :] = point[:, 0] + delta[2, :] # x2 delta[3, :] = point[:, 1] + delta[3, :] # y2 delta[0, :], delta[1, :], delta[2, :], delta[3, :] = corner2center(delta) return delta def _convert_score_numpy(self, score): def sofmax(logits): e_x = np.exp(logits) probs = e_x / np.sum(e_x, axis=-1, keepdims=True) return probs # score = score.permute(1, 2, 3, 0).contiguous().view(self.cls_out_channels, -1).permute(1, 0) # score = score.softmax(1).detach()[:, 1].cpu().numpy() score = score.transpose((1,2,3,0)).reshape(self.cls_out_channels, -1).transpose((1,0)) score = sofmax(score)[:,1] return score def _bbox_clip(self, cx, cy, width, height, boundary): cx = max(0, min(cx, boundary[1])) cy = max(0, min(cy, boundary[0])) width = max(10, min(width, boundary[1])) height = max(10, min(height, boundary[0])) return cx, cy, width, height def get_subwindow(self, im, pos, model_sz, original_sz, avg_chans): """ args: im: bgr based image pos: center position model_sz: exemplar size s_z: original size avg_chans: channel average """ if isinstance(pos, float): pos = [pos, pos] sz = original_sz im_sz = im.shape c = (original_sz + 1) / 2 # context_xmin = round(pos[0] - c) # py2 and py3 round context_xmin = np.floor(pos[0] - c + 0.5) context_xmax = context_xmin + sz - 1 # context_ymin = round(pos[1] - c) context_ymin = np.floor(pos[1] - c + 0.5) context_ymax = context_ymin + sz - 1 left_pad = int(max(0., -context_xmin)) top_pad = int(max(0., -context_ymin)) right_pad = int(max(0., context_xmax - im_sz[1] + 1)) bottom_pad = int(max(0., context_ymax - im_sz[0] + 1)) context_xmin = context_xmin + left_pad context_xmax = context_xmax + left_pad context_ymin = context_ymin + top_pad context_ymax = context_ymax + top_pad r, c, k = im.shape if any([top_pad, bottom_pad, left_pad, right_pad]): size = (r + top_pad + bottom_pad, c + left_pad + right_pad, k) te_im = np.zeros(size, np.uint8) te_im[top_pad:top_pad + r, left_pad:left_pad + c, :] = im if top_pad: te_im[0:top_pad, left_pad:left_pad + c, :] = avg_chans if bottom_pad: te_im[r + top_pad:, left_pad:left_pad + c, :] = avg_chans if left_pad: te_im[:, 0:left_pad, :] = avg_chans if right_pad: te_im[:, c + left_pad:, :] = avg_chans im_patch = te_im[int(context_ymin):int(context_ymax + 1), int(context_xmin):int(context_xmax + 1), :] else: im_patch = im[int(context_ymin):int(context_ymax + 1), int(context_xmin):int(context_xmax + 1), :] if not np.array_equal(model_sz, original_sz): im_patch = cv2.resize(im_patch, (model_sz, model_sz)) im_patch = im_patch.transpose(2, 0, 1) im_patch = im_patch[np.newaxis, :, :, :] im_patch = im_patch.astype(np.float32) return im_patch def init(self, img, bbox): """ args: img(np.ndarray): BGR image bbox: (x, y, w, h) bbox """ self.center_pos = np.array([bbox[0] + (bbox[2] - 1) / 2, bbox[1] + (bbox[3] - 1) / 2]) self.size = np.array([bbox[2], bbox[3]]) # calculate z crop size w_z = self.size[0