diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/__pycache__/__init__.cpython-39.pyc b/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..2214f20 Binary files /dev/null and b/__pycache__/__init__.cpython-39.pyc differ diff --git a/config/001_config.yaml b/config/config.yaml similarity index 100% rename from config/001_config.yaml rename to config/config.yaml diff --git a/docs/requirements b/docs/requirements index ffd3375..ff54449 100644 --- a/docs/requirements +++ b/docs/requirements @@ -17,6 +17,11 @@ │ ├── utils/ │ └── 006_utils.py # 工具函数 +├── tests/ # 测试各个程序 +│ ├── __init__.py +│ ├── 008_test_person_detector.py +│ ├── 009_test_distance_estimator.py +│ └── 010_test_api_server.py │ ├── 007_main.py # 主程序 └── requirements.txt # 项目依赖 \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..9ed6263 --- /dev/null +++ b/main.py @@ -0,0 +1,47 @@ +import yaml +import uvicorn +from src. 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/output/detection_20250113_110502_763716.jpg b/output/detection_20250113_110502_763716.jpg new file mode 100644 index 0000000..8237005 Binary files /dev/null and b/output/detection_20250113_110502_763716.jpg differ diff --git a/src/003_person_detector.py b/src/003_person_detector.py deleted file mode 100644 index edd2f9e..0000000 --- a/src/003_person_detector.py +++ /dev/null @@ -1,23 +0,0 @@ -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/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/__pycache__/__init__.cpython-39.pyc b/src/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..f969e31 Binary files /dev/null and b/src/__pycache__/__init__.cpython-39.pyc differ diff --git a/src/__pycache__/distance_estimator.cpython-39.pyc b/src/__pycache__/distance_estimator.cpython-39.pyc new file mode 100644 index 0000000..1df163c Binary files /dev/null and b/src/__pycache__/distance_estimator.cpython-39.pyc differ diff --git a/src/__pycache__/person_detector.cpython-39.pyc b/src/__pycache__/person_detector.cpython-39.pyc new file mode 100644 index 0000000..51c8340 Binary files /dev/null and b/src/__pycache__/person_detector.cpython-39.pyc differ diff --git a/src/__pycache__/ttt.cpython-39.pyc b/src/__pycache__/ttt.cpython-39.pyc new file mode 100644 index 0000000..2af70dd Binary files /dev/null and b/src/__pycache__/ttt.cpython-39.pyc differ diff --git a/src/005_api_server.py b/src/api_server.py similarity index 100% rename from src/005_api_server.py rename to src/api_server.py diff --git a/src/002_camera_handler.py b/src/camera_handler.py similarity index 100% rename from src/002_camera_handler.py rename to src/camera_handler.py diff --git a/src/004_distance_estimator.py b/src/distance_estimator.py similarity index 100% rename from src/004_distance_estimator.py rename to src/distance_estimator.py diff --git a/src/person_detector.py b/src/person_detector.py new file mode 100644 index 0000000..8987617 --- /dev/null +++ b/src/person_detector.py @@ -0,0 +1,51 @@ +from ultralytics import YOLO +import numpy as np +from ..utils.utils import VisualizationUtils + +class PersonDetector: + def __init__(self, model_path, conf_threshold=0.5): + self.model = YOLO(model_path) + self.conf_threshold = conf_threshold + self.vis_utils = VisualizationUtils() + + def detect(self, frame, save_visualization=False, save_path="output"): + """ + 检测图像中的人物并可选择保存可视化结果 + + Args: + frame: 输入图像 + save_visualization: 是否保存可视化结果 + save_path: 可视化结果保存路径 + + Returns: + persons: 检测到的人物列表 + vis_path: 如果save_visualization为True,返回保存的文件路径,否则为None + """ + 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) + }) + + vis_path = None + if save_visualization and persons: + # 绘制检测结果 + vis_image = self.vis_utils.draw_detection_results( + frame, + persons, + show_confidence=True + ) + # 保存结果 + vis_path = self.vis_utils.save_detection_image( + vis_image, + save_path + ) + + return persons, vis_path \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/__pycache__/008_test_person_detector.cpython-39-pytest-8.3.4.pyc b/tests/__pycache__/008_test_person_detector.cpython-39-pytest-8.3.4.pyc new file mode 100644 index 0000000..86bd3a8 Binary files /dev/null and b/tests/__pycache__/008_test_person_detector.cpython-39-pytest-8.3.4.pyc differ diff --git a/tests/__pycache__/__init__.cpython-39.pyc b/tests/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..c110a67 Binary files /dev/null and b/tests/__pycache__/__init__.cpython-39.pyc differ diff --git a/tests/__pycache__/test_distance_estimator.cpython-39-pytest-8.3.4.pyc b/tests/__pycache__/test_distance_estimator.cpython-39-pytest-8.3.4.pyc new file mode 100644 index 0000000..6ed2aa8 Binary files /dev/null and b/tests/__pycache__/test_distance_estimator.cpython-39-pytest-8.3.4.pyc differ diff --git a/tests/__pycache__/test_person_detector.cpython-39-pytest-8.3.4.pyc b/tests/__pycache__/test_person_detector.cpython-39-pytest-8.3.4.pyc new file mode 100644 index 0000000..7539477 Binary files /dev/null and b/tests/__pycache__/test_person_detector.cpython-39-pytest-8.3.4.pyc differ diff --git a/tests/test_api_server.py b/tests/test_api_server.py new file mode 100644 index 0000000..210888d --- /dev/null +++ b/tests/test_api_server.py @@ -0,0 +1,40 @@ +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel +from typing import List, Optional + +class Distance(BaseModel): + distance_mm: float + confidence: float + detection_image_path: Optional[str] = None + +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, vis_path = self.detector.detect( + frame, + save_visualization=True, + save_path="output/detections" + ) + + 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'], + detection_image_path=vis_path + )) + + return distances \ No newline at end of file diff --git a/tests/test_data/2024-04-02-04_14_55-749386_jpg.rf.6e256601cc6513bf8e6f4497f4662d71.jpg b/tests/test_data/2024-04-02-04_14_55-749386_jpg.rf.6e256601cc6513bf8e6f4497f4662d71.jpg new file mode 100644 index 0000000..bf3b4b5 Binary files /dev/null and b/tests/test_data/2024-04-02-04_14_55-749386_jpg.rf.6e256601cc6513bf8e6f4497f4662d71.jpg differ diff --git a/tests/test_data/Meeting-2nd-Floor-8-19-2024-13_05_00-GMT-7-8-19-2024-13_35_30-GMT-7_mp4-1803_jpg.rf.3e48dd21b87dfe87dd5c32df99b3c468.jpg b/tests/test_data/Meeting-2nd-Floor-8-19-2024-13_05_00-GMT-7-8-19-2024-13_35_30-GMT-7_mp4-1803_jpg.rf.3e48dd21b87dfe87dd5c32df99b3c468.jpg new file mode 100644 index 0000000..6ce44ac Binary files /dev/null and b/tests/test_data/Meeting-2nd-Floor-8-19-2024-13_05_00-GMT-7-8-19-2024-13_35_30-GMT-7_mp4-1803_jpg.rf.3e48dd21b87dfe87dd5c32df99b3c468.jpg differ diff --git a/tests/test_data/person.jpg b/tests/test_data/person.jpg new file mode 100644 index 0000000..5bb8e61 Binary files /dev/null and b/tests/test_data/person.jpg differ diff --git a/tests/test_distance_estimator.py b/tests/test_distance_estimator.py new file mode 100644 index 0000000..e0d91ca --- /dev/null +++ b/tests/test_distance_estimator.py @@ -0,0 +1,54 @@ +import pytest +import numpy as np +from ..src.distance_estimator import DistanceEstimator + +@pytest.fixture +def estimator(): + return DistanceEstimator( + focal_length=1000, # mm + sensor_height=5.6, # mm + avg_person_height=1700 # mm + ) + +def test_distance_estimator_initialization(estimator): + assert estimator.focal_length == 1000 + assert estimator.sensor_height == 5.6 + assert estimator.avg_person_height == 1700 + +def test_distance_estimation(estimator): + # 测试不同像素高度的情况 + test_cases = [ + # (person_height_pixels, frame_height, expected_distance_range) + (400, 720, (2000, 5000)), # 近距离 + (200, 720, (4000, 10000)), # 中等距离 + (100, 720, (8000, 20000)) # 远距离 + ] + + for height_px, frame_height, (min_dist, max_dist) in test_cases: + distance = estimator.estimate_distance(height_px, frame_height) + assert min_dist <= distance <= max_dist, \ + f"Distance {distance}mm not in expected range [{min_dist}, {max_dist}]" + +def test_distance_estimation_edge_cases(estimator): + # 测试边界情况 + with pytest.raises(ZeroDivisionError): + estimator.estimate_distance(0, 720) + + # 测试极小值 + distance = estimator.estimate_distance(1, 720) + assert distance > 0 + + # 测试当人物高度等于图像高度的情况 + distance = estimator.estimate_distance(720, 720) + assert distance > 0 + +def test_distance_estimation_linearity(estimator): + # 测试距离估算的线性关系 + height1, height2 = 200, 100 + frame_height = 720 + + distance1 = estimator.estimate_distance(height1, frame_height) + distance2 = estimator.estimate_distance(height2, frame_height) + + # 当像素高度减半时,距离应该翻倍 + assert abs(distance2 / distance1 - 2.0) < 0.1 \ No newline at end of file diff --git a/tests/test_person_detector.py b/tests/test_person_detector.py new file mode 100644 index 0000000..0899df5 --- /dev/null +++ b/tests/test_person_detector.py @@ -0,0 +1,51 @@ +import pytest +import numpy as np +import cv2 +# import os +# import sys + +from ..src.person_detector import PersonDetector +# from src.person_detector import PersonDetector + +@pytest.fixture +def detector(): + # 使用YOLOv8n模型进行测试 + return PersonDetector("yolov8n.pt", conf_threshold=0.5) + +@pytest.fixture +def sample_image(): + # 创建一个测试用的图像,包含一个人的图像 + img_path = "tests/test_data/person.jpg" # 请准备一张包含人物的测试图片 + return cv2.imread(img_path) + +def test_person_detector_initialization(detector): + assert detector.model is not None + assert detector.conf_threshold == 0.5 + +def test_person_detection(detector, sample_image): + persons = detector.detect(sample_image, True) + + # 验证检测结果 + assert isinstance(persons, list) + if len(persons) > 0: + person = persons[0] + assert 'bbox' in person + assert 'confidence' in person + + # 验证边界框格式 + bbox = person['bbox'] + assert len(bbox) == 4 + assert all(isinstance(x, int) for x in bbox) + assert bbox[0] < bbox[2] # x1 < x2 + assert bbox[1] < bbox[3] # y1 < y2 + + # 验证置信度 + print("人物置信度", person['confidence']) + assert 0 <= person['confidence'] <= 1 + +def test_person_detector_empty_image(detector): + # 测试空图像 + empty_image = np.zeros((100, 100, 3), dtype=np.uint8) + persons = detector.detect(empty_image) + assert isinstance(persons, list) + assert len(persons) == 0 \ No newline at end of file diff --git a/utils/__pycache__/utils.cpython-39.pyc b/utils/__pycache__/utils.cpython-39.pyc new file mode 100644 index 0000000..56ec7c1 Binary files /dev/null and b/utils/__pycache__/utils.cpython-39.pyc differ diff --git a/utils/utils.py b/utils/utils.py new file mode 100644 index 0000000..fbb069e --- /dev/null +++ b/utils/utils.py @@ -0,0 +1,117 @@ +import cv2 +import numpy as np +from datetime import datetime +import os + +class VisualizationUtils: + @staticmethod + def draw_detection_results( + image: np.ndarray, + persons: list, + distances: list = None, + save_path: str = "output", + show_confidence: bool = True + ) -> np.ndarray: + """ + 在图像上绘制检测结果和距离信息 + + Args: + image: 原始图像 + persons: 检测到的人物列表,每个元素包含bbox和confidence + distances: 对应的距离列表(可选) + save_path: 保存路径 + show_confidence: 是否显示置信度 + + Returns: + 绘制了检测结果的图像 + """ + # 创建图像副本 + vis_image = image.copy() + + # 定义颜色和字体 + BOX_COLOR = (0, 255, 0) # 绿色边框 + TEXT_COLOR = (255, 255, 255) # 白色文字 + FONT = cv2.FONT_HERSHEY_SIMPLEX + FONT_SCALE = 0.6 + THICKNESS = 2 + + # 遍历每个检测结果 + for idx, person in enumerate(persons): + bbox = person['bbox'] + conf = person['confidence'] + + # 绘制边界框 + cv2.rectangle( + vis_image, + (bbox[0], bbox[1]), + (bbox[2], bbox[3]), + BOX_COLOR, + THICKNESS + ) + + # 准备显示文本 + text_items = [] + if show_confidence: + text_items.append(f"Conf: {conf:.2f}") + + if distances and idx < len(distances): + dist_meters = distances[idx] / 1000 # 转换为米 + text_items.append(f"Dist: {dist_meters:.2f}m") + + # 绘制文本 + text = " | ".join(text_items) + if text: + # 获取文本大小 + (text_width, text_height), _ = cv2.getTextSize( + text, FONT, FONT_SCALE, THICKNESS + ) + + # 绘制文本背景 + cv2.rectangle( + vis_image, + (bbox[0], bbox[1] - text_height - 10), + (bbox[0] + text_width + 10, bbox[1]), + BOX_COLOR, + -1 # 填充矩形 + ) + + # 绘制文本 + cv2.putText( + vis_image, + text, + (bbox[0] + 5, bbox[1] - 5), + FONT, + FONT_SCALE, + TEXT_COLOR, + THICKNESS + ) + + return vis_image + + @staticmethod + def save_detection_image( + image: np.ndarray, + save_path: str = "output" + ) -> str: + """ + 保存检测结果图像 + + Args: + image: 要保存的图像 + save_path: 保存路径 + + Returns: + 保存的文件路径 + """ + # 确保输出目录存在 + os.makedirs(save_path, exist_ok=True) + + # 生成文件名(使用时间戳) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + filename = f"detection_{timestamp}.jpg" + filepath = os.path.join(save_path, filename) + + # 保存图像 + cv2.imwrite(filepath, image) + + return filepath \ No newline at end of file diff --git a/yolov8n.pt b/yolov8n.pt new file mode 100644 index 0000000..0db4ca4 Binary files /dev/null and b/yolov8n.pt differ