first commit

007_main.py

@@ -0,0 +1,47 @@
+import yaml
+import uvicorn
+from src.camera_handler import RTSPCamera
+from src.person_detector import PersonDetector
+from src.distance_estimator import DistanceEstimator
+from src.api_server import DistanceAPI, app
+
+def load_config():
+    with open('config/config.yaml', 'r') as f:
+        return yaml.safe_load(f)
+
+def main():
+    config = load_config()
+    
+    # 初始化摄像头
+    camera = RTSPCamera(
+        config['camera']['rtsp_url'],
+        config['camera']['fps']
+    )
+    camera.start()
+    
+    # 初始化检测器
+    detector = PersonDetector(
+        config['model']['person_detection']['model_name'],
+        config['model']['person_detection']['confidence_threshold']
+    )
+    
+    # 初始化距离估算器
+    estimator = DistanceEstimator(
+        config['model']['distance_estimation']['focal_length'],
+        config['model']['distance_estimation']['sensor_height'],
+        config['model']['distance_estimation']['average_person_height']
+    )
+    
+    # 初始化API
+    distance_api = DistanceAPI(camera, detector, estimator)
+    app.distance_api = distance_api
+    
+    # 启动API服务器
+    uvicorn.run(
+        app,
+        host=config['api']['host'],
+        port=config['api']['port']
+    )
+
+if __name__ == "__main__":
+    main()

config/001_config.yaml

@@ -0,0 +1,19 @@
+camera:
+  rtsp_url: "rtsp://username:password@ip:port/stream"
+  fps: 30
+  width: 1280
+  height: 720
+
+model:
+  person_detection:
+    model_name: "yolov8n.pt"  # 使用YOLOv8进行人物检测
+    confidence_threshold: 0.5
+    
+  distance_estimation:
+    focal_length: 1000  # 摄像头焦距
+    sensor_height: 5.6  # 摄像头传感器高度(mm)
+    average_person_height: 1700  # 平均人身高(mm)
+
+api:
+  host: "0.0.0.0"
+  port: 8000

docs/requirements

@@ -0,0 +1,22 @@
+项目需求:
+    1.实时读取摄像头中RTSP流中的画面并处理.
+    2.对于摄像头的画面找出其中的人物,并估计摄像头到人物的距离.
+    3.写一个接口.调用接口时,返回当前摄像头中人物到摄像头的距离.
+
+项目结构:
+/
+│
+├── config/
+│   └── 001_config.yaml           # 配置文件
+│
+├── src/
+│   ├── 002_camera_handler.py     # 摄像头处理模块
+│   ├── 003_person_detector.py    # 人物检测模块
+│   ├── 004_distance_estimator.py # 距离估计模块
+│   └── 005_api_server.py         # API接口服务
+│
+├── utils/
+│   └── 006_utils.py              # 工具函数
+│
+├── 007_main.py                   # 主程序
+└── requirements.txt              # 项目依赖

requirements.txt

@@ -0,0 +1,6 @@
+opencv-python>=4.5.0
+numpy>=1.19.0
+ultralytics>=8.0.0
+fastapi>=0.68.0
+uvicorn>=0.15.0
+pyyaml>=5.4.0

src/002_camera_handler.py

@@ -0,0 +1,34 @@
+import cv2
+import numpy as np
+from threading import Thread
+import time
+
+# 摄像头处理模块
+class RTSPCamera:
+    def __init__(self, rtsp_url, fps=30):
+        self.rtsp_url = rtsp_url
+        self.fps = fps
+        self.frame = None
+        self.running = False
+        self._thread = None
+        
+    def start(self):
+        self.running = True
+        self._thread = Thread(target=self._update_frame, daemon=True)
+        self._thread.start()
+        
+    def _update_frame(self):
+        cap = cv2.VideoCapture(self.rtsp_url)
+        while self.running:
+            ret, frame = cap.read()
+            if ret:
+                self.frame = frame
+            time.sleep(1/self.fps)
+            
+    def get_frame(self):
+        return self.frame.copy() if self.frame is not None else None
+    
+    def stop(self):
+        self.running = False
+        if self._thread is not None:
+            self._thread.join()

src/003_person_detector.py

@@ -0,0 +1,23 @@
+from ultralytics import YOLO
+import numpy as np
+
+class PersonDetector:
+    def __init__(self, model_path, conf_threshold=0.5):
+        self.model = YOLO(model_path)
+        self.conf_threshold = conf_threshold
+        
+    def detect(self, frame):
+        results = self.model(frame, classes=0)  # 只检测人类
+        persons = []
+        
+        for result in results:
+            boxes = result.boxes
+            for box in boxes:
+                if box.conf > self.conf_threshold:
+                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+                    persons.append({
+                        'bbox': (int(x1), int(y1), int(x2), int(y2)),
+                        'confidence': float(box.conf)
+                    })
+        
+        return persons

src/004_distance_estimator.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+class DistanceEstimator:
+    def __init__(self, focal_length, sensor_height, avg_person_height):
+        self.focal_length = focal_length
+        self.sensor_height = sensor_height
+        self.avg_person_height = avg_person_height
+        
+    def estimate_distance(self, person_height_pixels, frame_height):
+        # 使用相似三角形原理估算距离
+        real_height = self.avg_person_height  # mm
+        sensor_height_pixels = frame_height
+        
+        distance = (real_height * self.focal_length) / (person_height_pixels * self.sensor_height / sensor_height_pixels)
+        return distance  # 返回距离(mm)

src/005_api_server.py

@@ -0,0 +1,35 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
+
+app = FastAPI()
+
+class Distance(BaseModel):
+    distance_mm: float
+    confidence: float
+
+class DistanceAPI:
+    def __init__(self, camera, detector, estimator):
+        self.camera = camera
+        self.detector = detector
+        self.estimator = estimator
+        
+    async def get_distances(self) -> List[Distance]:
+        frame = self.camera.get_frame()
+        if frame is None:
+            raise HTTPException(status_code=500, detail="No frame available")
+            
+        persons = self.detector.detect(frame)
+        distances = []
+        
+        for person in persons:
+            bbox = person['bbox']
+            person_height = bbox[3] - bbox[1]  # y2 - y1
+            distance = self.estimator.estimate_distance(person_height, frame.shape[0])
+            
+            distances.append(Distance(
+                distance_mm=distance,
+                confidence=person['confidence']
+            ))
+            
+        return distances