|
| 1 | +import os |
| 2 | +import cv2 |
| 3 | +import sys |
| 4 | +import argparse |
| 5 | +import time |
| 6 | +import numpy as np |
| 7 | +import threading |
| 8 | +from fastapi import FastAPI, Response |
| 9 | +from fastapi.responses import StreamingResponse |
| 10 | +import uvicorn |
| 11 | + |
| 12 | +# 导入共享工具 |
| 13 | +from py_utils.coco_utils import COCO_test_helper |
| 14 | + |
| 15 | +# 尝试导入RKNN-Toolkit-Lite2 |
| 16 | +try: |
| 17 | + from rknnlite.api import RKNNLite |
| 18 | + RKNN_LITE_AVAILABLE = True |
| 19 | +except ImportError: |
| 20 | + RKNN_LITE_AVAILABLE = False |
| 21 | + print("Warning: RKNN-Toolkit-Lite2 not available, using fallback") |
| 22 | + |
| 23 | +# 常量定义 |
| 24 | +OBJ_THRESH = 0.25 |
| 25 | +NMS_THRESH = 0.45 |
| 26 | +IMG_SIZE = (640, 640) # (width, height) |
| 27 | + |
| 28 | +CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light", |
| 29 | + "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant", |
| 30 | + "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite", |
| 31 | + "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ", |
| 32 | + "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa", |
| 33 | + "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop ","mouse ","remote ","keyboard ","cell phone","microwave ", |
| 34 | + "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush ") |
| 35 | + |
| 36 | +# --- FastAPI 核心组件 --- |
| 37 | +app = FastAPI(title="reComputer RK-CV Web Preview (RK3576)") |
| 38 | + |
| 39 | +class FrameBuffer: |
| 40 | + def __init__(self): |
| 41 | + self.frame = None |
| 42 | + self.lock = threading.Lock() |
| 43 | + |
| 44 | + def set_frame(self, frame): |
| 45 | + with self.lock: |
| 46 | + self.frame = frame |
| 47 | + |
| 48 | + def get_frame(self): |
| 49 | + with self.lock: |
| 50 | + return self.frame |
| 51 | + |
| 52 | +frame_buffer = FrameBuffer() |
| 53 | + |
| 54 | +@app.get("/api/video_feed") |
| 55 | +async def video_feed(): |
| 56 | + def generate(): |
| 57 | + while True: |
| 58 | + frame = frame_buffer.get_frame() |
| 59 | + if frame is not None: |
| 60 | + yield (b'--frame\r\n' |
| 61 | + b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n') |
| 62 | + time.sleep(0.03) |
| 63 | + return StreamingResponse(generate(), media_type="multipart/x-mixed-replace; boundary=frame") |
| 64 | + |
| 65 | +@app.get("/") |
| 66 | +async def index(): |
| 67 | + return Response(content=""" |
| 68 | + <html> |
| 69 | + <head><title>reComputer RK-CV Web Preview</title></head> |
| 70 | + <body style="background-color: #1a1a1a; color: white; text-align: center; font-family: sans-serif;"> |
| 71 | + <h1>reComputer RK-CV Real-time Detection (RK3576 Web Mode)</h1> |
| 72 | + <div style="margin: 20px auto; display: inline-block; border: 5px solid #333; border-radius: 10px; overflow: hidden;"> |
| 73 | + <img src="/api/video_feed" style="max-width: 100%; height: auto;"> |
| 74 | + </div> |
| 75 | + <p>Streaming via FastAPI + MJPEG | Port: 8000</p> |
| 76 | + </body> |
| 77 | + </html> |
| 78 | + """, media_type="text/html") |
| 79 | + |
| 80 | +def run_fastapi(host, port): |
| 81 | + uvicorn.run(app, host=host, port=port, log_level="error") |
| 82 | + |
| 83 | +# --- 推理逻辑 --- |
| 84 | + |
| 85 | +def filter_boxes(boxes, box_confidences, box_class_probs): |
| 86 | + """Filter boxes with object threshold.""" |
| 87 | + box_confidences = box_confidences.reshape(-1) |
| 88 | + class_max_score = np.max(box_class_probs, axis=-1) |
| 89 | + classes = np.argmax(box_class_probs, axis=-1) |
| 90 | + _class_pos = np.where(class_max_score* box_confidences >= OBJ_THRESH) |
| 91 | + scores = (class_max_score* box_confidences)[_class_pos] |
| 92 | + boxes = boxes[_class_pos] |
| 93 | + classes = classes[_class_pos] |
| 94 | + return boxes, classes, scores |
| 95 | + |
| 96 | +def nms_boxes(boxes, scores): |
| 97 | + """Suppress non-maximal boxes.""" |
| 98 | + x = boxes[:, 0] |
| 99 | + y = boxes[:, 1] |
| 100 | + w = boxes[:, 2] - boxes[:, 0] |
| 101 | + h = boxes[:, 3] - boxes[:, 1] |
| 102 | + areas = w * h |
| 103 | + order = scores.argsort()[::-1] |
| 104 | + keep = [] |
| 105 | + while order.size > 0: |
| 106 | + i = order[0] |
| 107 | + keep.append(i) |
| 108 | + xx1 = np.maximum(x[i], x[order[1:]]) |
| 109 | + yy1 = np.maximum(y[i], y[order[1:]]) |
| 110 | + xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) |
| 111 | + yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) |
| 112 | + w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) |
| 113 | + h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) |
| 114 | + inter = w1 * h1 |
| 115 | + ovr = inter / (areas[i] + areas[order[1:]] - inter) |
| 116 | + inds = np.where(ovr <= NMS_THRESH)[0] |
| 117 | + order = order[inds + 1] |
| 118 | + return np.array(keep) |
| 119 | + |
| 120 | +def dfl(position): |
| 121 | + n, c, h, w = position.shape |
| 122 | + p_num = 4 |
| 123 | + mc = c // p_num |
| 124 | + y = position.reshape(n, p_num, mc, h, w) |
| 125 | + y_exp = np.exp(y - np.max(y, axis=2, keepdims=True)) |
| 126 | + y_softmax = y_exp / np.sum(y_exp, axis=2, keepdims=True) |
| 127 | + acc_metrix = np.arange(mc).reshape(1, 1, mc, 1, 1).astype(np.float32) |
| 128 | + y = (y_softmax * acc_metrix).sum(2) |
| 129 | + return y |
| 130 | + |
| 131 | +def box_process(position): |
| 132 | + grid_h, grid_w = position.shape[2:4] |
| 133 | + col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h)) |
| 134 | + col = col.reshape(1, 1, grid_h, grid_w) |
| 135 | + row = row.reshape(1, 1, grid_h, grid_w) |
| 136 | + grid = np.concatenate((col, row), axis=1) |
| 137 | + stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1) |
| 138 | + position = dfl(position) |
| 139 | + box_xy = grid +0.5 -position[:,0:2,:,:] |
| 140 | + box_xy2 = grid +0.5 +position[:,2:4,:,:] |
| 141 | + xyxy = np.concatenate((box_xy*stride, box_xy2*stride), axis=1) |
| 142 | + return xyxy |
| 143 | + |
| 144 | +def post_process(input_data): |
| 145 | + if input_data is None: |
| 146 | + return None, None, None |
| 147 | + boxes, scores, classes_conf = [], [], [] |
| 148 | + defualt_branch=3 |
| 149 | + pair_per_branch = len(input_data)//defualt_branch |
| 150 | + for i in range(defualt_branch): |
| 151 | + boxes.append(box_process(input_data[pair_per_branch*i])) |
| 152 | + classes_conf.append(input_data[pair_per_branch*i+1]) |
| 153 | + scores.append(np.ones_like(input_data[pair_per_branch*i+1][:,:1,:,:], dtype=np.float32)) |
| 154 | + |
| 155 | + def sp_flatten(_in): |
| 156 | + ch = _in.shape[1] |
| 157 | + _in = _in.transpose(0,2,3,1) |
| 158 | + return _in.reshape(-1, ch) |
| 159 | + |
| 160 | + boxes = [sp_flatten(_v) for _v in boxes] |
| 161 | + classes_conf = [sp_flatten(_v) for _v in classes_conf] |
| 162 | + scores = [sp_flatten(_v) for _v in scores] |
| 163 | + boxes = np.concatenate(boxes) |
| 164 | + classes_conf = np.concatenate(classes_conf) |
| 165 | + scores = np.concatenate(scores) |
| 166 | + boxes, classes, scores = filter_boxes(boxes, scores, classes_conf) |
| 167 | + nboxes, nclasses, nscores = [], [], [] |
| 168 | + for c in set(classes): |
| 169 | + inds = np.where(classes == c) |
| 170 | + b = boxes[inds] |
| 171 | + c = classes[inds] |
| 172 | + s = scores[inds] |
| 173 | + keep = nms_boxes(b, s) |
| 174 | + if len(keep) != 0: |
| 175 | + nboxes.append(b[keep]) |
| 176 | + nclasses.append(c[keep]) |
| 177 | + nscores.append(s[keep]) |
| 178 | + if not nclasses and not nscores: |
| 179 | + return None, None, None |
| 180 | + boxes = np.concatenate(nboxes) |
| 181 | + classes = np.concatenate(nclasses) |
| 182 | + scores = np.concatenate(nscores) |
| 183 | + return boxes, classes, scores |
| 184 | + |
| 185 | +def draw(image, boxes, scores, classes): |
| 186 | + for box, score, cl in zip(boxes, scores, classes): |
| 187 | + top, left, right, bottom = [int(_b) for _b in box] |
| 188 | + cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2) |
| 189 | + cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score), |
| 190 | + (top, left - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) |
| 191 | + |
| 192 | +class RKNNLiteModel: |
| 193 | + def __init__(self, model_path): |
| 194 | + if not RKNN_LITE_AVAILABLE: |
| 195 | + raise ImportError("RKNN-Toolkit-Lite2 is not available") |
| 196 | + if not os.path.exists(model_path): |
| 197 | + raise FileNotFoundError(f"RKNN model file not found: {model_path}") |
| 198 | + self.rknn_lite = RKNNLite() |
| 199 | + print(f'Loading RKNN model from {model_path}...') |
| 200 | + ret = self.rknn_lite.load_rknn(model_path) |
| 201 | + if ret != 0: |
| 202 | + raise Exception(f"Load RKNN model failed with error code: {ret}") |
| 203 | + print('Initializing runtime...') |
| 204 | + ret = self.rknn_lite.init_runtime() |
| 205 | + if ret != 0: |
| 206 | + raise Exception(f"Init runtime failed with error code: {ret}") |
| 207 | + print('RKNN model loaded successfully') |
| 208 | + |
| 209 | + def run(self, inputs): |
| 210 | + try: |
| 211 | + if len(inputs.shape) == 3: |
| 212 | + inputs = np.expand_dims(inputs, axis=0) |
| 213 | + if inputs.dtype != np.uint8: |
| 214 | + inputs = inputs.astype(np.uint8) |
| 215 | + return self.rknn_lite.inference(inputs=[inputs]) |
| 216 | + except Exception as e: |
| 217 | + print(f"Inference error: {e}") |
| 218 | + return None |
| 219 | + |
| 220 | + def release(self): |
| 221 | + if hasattr(self, 'rknn_lite'): |
| 222 | + self.rknn_lite.release() |
| 223 | + |
| 224 | +def preprocess_frame(frame, co_helper): |
| 225 | + img = co_helper.letter_box(im=frame.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0]), pad_color=(0,0,0)) |
| 226 | + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
| 227 | + return img |
| 228 | + |
| 229 | +def main(): |
| 230 | + parser = argparse.ArgumentParser(description='Object detection on RK3576 (Web Preview Mode)') |
| 231 | + parser.add_argument('--model_path', type=str, required=True, help='RKNN model path') |
| 232 | + parser.add_argument('--camera_id', type=int, default=1, help='Camera device ID (default: 1 for /dev/video1)') |
| 233 | + parser.add_argument('--video_path', type=str, help='Path to video file (overrides camera_id)') |
| 234 | + parser.add_argument('--host', type=str, default='0.0.0.0', help='Web server host') |
| 235 | + parser.add_argument('--port', type=int, default=8000, help='Web server port') |
| 236 | + args = parser.parse_args() |
| 237 | + |
| 238 | + if not RKNN_LITE_AVAILABLE: |
| 239 | + print("Error: RKNN-Toolkit-Lite2 is not available.") |
| 240 | + return |
| 241 | + |
| 242 | + # 启动 Web 服务器线程 |
| 243 | + web_thread = threading.Thread(target=run_fastapi, args=(args.host, args.port), daemon=True) |
| 244 | + web_thread.start() |
| 245 | + print(f"Web Preview started at http://{args.host}:{args.port}") |
| 246 | + |
| 247 | + # 初始化模型 |
| 248 | + model = RKNNLiteModel(args.model_path) |
| 249 | + co_helper = COCO_test_helper(enable_letter_box=True) |
| 250 | + |
| 251 | + # 打开视频源 |
| 252 | + if args.video_path: |
| 253 | + cap = cv2.VideoCapture(args.video_path) |
| 254 | + else: |
| 255 | + cap = cv2.VideoCapture(args.camera_id) |
| 256 | + cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) |
| 257 | + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) |
| 258 | + |
| 259 | + if not cap.isOpened(): |
| 260 | + print(f"Error: Cannot open video source (ID: {args.camera_id if not args.video_path else args.video_path})") |
| 261 | + return |
| 262 | + |
| 263 | + fps_counter = 0 |
| 264 | + try: |
| 265 | + while True: |
| 266 | + ret, frame = cap.read() |
| 267 | + if not ret: |
| 268 | + if args.video_path: |
| 269 | + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) |
| 270 | + continue |
| 271 | + break |
| 272 | + |
| 273 | + # 推理流程 |
| 274 | + processed_img = preprocess_frame(frame, co_helper) |
| 275 | + start_time = time.time() |
| 276 | + outputs = model.run(processed_img) |
| 277 | + inference_time = time.time() - start_time |
| 278 | + |
| 279 | + if outputs is not None: |
| 280 | + boxes, classes, scores = post_process(outputs) |
| 281 | + if boxes is not None: |
| 282 | + draw(frame, co_helper.get_real_box(boxes), scores, classes) |
| 283 | + |
| 284 | + # 计算并显示 FPS |
| 285 | + inf_fps = 1.0 / inference_time if inference_time > 0 else 0 |
| 286 | + fps_counter = 0.9 * fps_counter + 0.1 * inf_fps if fps_counter > 0 else inf_fps |
| 287 | + cv2.putText(frame, f'NPU FPS: {fps_counter:.1f}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) |
| 288 | + |
| 289 | + # 更新 Web 帧缓冲区 |
| 290 | + _, buffer = cv2.imencode('.jpg', frame) |
| 291 | + frame_buffer.set_frame(buffer.tobytes()) |
| 292 | + |
| 293 | + # 降低 CPU 占用 |
| 294 | + time.sleep(0.01) |
| 295 | + |
| 296 | + except KeyboardInterrupt: |
| 297 | + print("Interrupted by user") |
| 298 | + finally: |
| 299 | + cap.release() |
| 300 | + model.release() |
| 301 | + |
| 302 | +if __name__ == '__main__': |
| 303 | + main() |
0 commit comments