From 336ed3dadcd35d0d80de4a4a4706f6f18a31ea41 Mon Sep 17 00:00:00 2001
From: haotian <2421912570@qq.com>
Date: Mon, 13 Jan 2025 10:22:55 +0800
Subject: [PATCH] first commit

---
 007_main.py                   | 47 +++++++++++++++++++++++++++++++++++
 config/001_config.yaml        | 19 ++++++++++++++
 docs/requirements             | 22 ++++++++++++++++
 requirements.txt              |  6 +++++
 src/002_camera_handler.py     | 34 +++++++++++++++++++++++++
 src/003_person_detector.py    | 23 +++++++++++++++++
 src/004_distance_estimator.py | 15 +++++++++++
 src/005_api_server.py         | 35 ++++++++++++++++++++++++++
 8 files changed, 201 insertions(+)
 create mode 100644 007_main.py
 create mode 100644 config/001_config.yaml
 create mode 100644 docs/requirements
 create mode 100644 requirements.txt
 create mode 100644 src/002_camera_handler.py
 create mode 100644 src/003_person_detector.py
 create mode 100644 src/004_distance_estimator.py
 create mode 100644 src/005_api_server.py

diff --git a/007_main.py b/007_main.py
new file mode 100644
index 0000000..1ace5b4
--- /dev/null
+++ b/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()
\ No newline at end of file
diff --git a/config/001_config.yaml b/config/001_config.yaml
new file mode 100644
index 0000000..7b1d3a7
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/docs/requirements b/docs/requirements
new file mode 100644
index 0000000..ffd3375
--- /dev/null
+++ b/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              # 项目依赖
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..14ffc9d
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/src/002_camera_handler.py b/src/002_camera_handler.py
new file mode 100644
index 0000000..18a3f15
--- /dev/null
+++ b/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()
\ No newline at end of file
diff --git a/src/003_person_detector.py b/src/003_person_detector.py
new file mode 100644
index 0000000..edd2f9e
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/src/004_distance_estimator.py b/src/004_distance_estimator.py
new file mode 100644
index 0000000..558bd72
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/src/005_api_server.py b/src/005_api_server.py
new file mode 100644
index 0000000..32eb9e3
--- /dev/null
+++ b/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
\ No newline at end of file