camera.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. import cv2
  2. import time
  3. from datetime import datetime
  4. import os
  5. import uuid
  6. from fastapi import APIRouter, HTTPException
  7. from pydantic import BaseModel
  8. from concurrent.futures import ThreadPoolExecutor
  9. from threading import Lock, Timer
  10. from app.core.config import settings
  11. # -------------------- 配置 --------------------
  12. executor = ThreadPoolExecutor(max_workers=4)
  13. lock = Lock()
  14. # 摄像头任务状态
  15. cam_status = {}
  16. router = APIRouter()
  17. class TaskRequest(BaseModel):
  18. camera_id: str
  19. task_name: str
  20. # -------------------- 录像函数 --------------------
  21. def record_camera(camera_id: str, task_name: str, output_file: str):
  22. rtsp_url = settings.CAMERA_CONFIG[camera_id]
  23. cap = cv2.VideoCapture(rtsp_url)
  24. cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) # 减少延迟
  25. if not cap.isOpened():
  26. raise RuntimeError(f"无法连接摄像头 {camera_id}")
  27. frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  28. frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  29. fps = cap.get(cv2.CAP_PROP_FPS)
  30. if fps <= 0: fps = 25
  31. fourcc = cv2.VideoWriter_fourcc(*'mp4v')
  32. out = cv2.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height))
  33. print(f"[{camera_id}] 任务 {task_name} 开始录制...")
  34. try:
  35. while True:
  36. ret, frame = cap.read()
  37. if not ret:
  38. print(f"[{camera_id}] 丢帧,停止录像")
  39. break
  40. out.write(frame)
  41. # 检查任务是否被标记停止
  42. with lock:
  43. status = cam_status.get(camera_id)
  44. if not status or status.get("stop_flag"):
  45. break
  46. time.sleep(0.001)
  47. finally:
  48. cap.release()
  49. out.release()
  50. print(f"[{camera_id}] 任务 {task_name} 完成,文件: {output_file}")
  51. def stop_task_internal(camera_id: str):
  52. with lock:
  53. status = cam_status.get(camera_id)
  54. if not status:
  55. return
  56. status["stop_flag"] = True
  57. status["timer"].cancel()
  58. # -------------------- API接口 --------------------
  59. @router.post("/start_task")
  60. def start_task(req: TaskRequest):
  61. camera_id = req.camera_id
  62. task_name = req.task_name
  63. if camera_id not in settings.CAMERA_CONFIG:
  64. raise HTTPException(status_code=404, detail=f"摄像头 {camera_id} 未配置")
  65. with lock:
  66. if camera_id in cam_status:
  67. return {"status": "running", "task_name": cam_status[camera_id]["task_name"]}
  68. # 输出文件名
  69. now_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
  70. filename = f"{camera_id}_{task_name}_{now_time}.mp4"
  71. output_file = os.path.join(settings.OUTPUT_DIR, filename)
  72. # 记录状态
  73. stop_flag = False
  74. timer = Timer(settings.MAX_TASK_SECONDS, lambda: stop_task_internal(camera_id))
  75. timer.start()
  76. future = executor.submit(record_camera, camera_id, task_name, output_file)
  77. cam_status[camera_id] = {
  78. "task_name": task_name,
  79. "start_time": time.time(),
  80. "timer": timer,
  81. "future": future,
  82. "file": output_file,
  83. "stop_flag": stop_flag
  84. }
  85. return {"status": "started", "file": output_file}
  86. @router.post("/stop_task")
  87. def stop_task(req: TaskRequest):
  88. camera_id = req.camera_id
  89. task_name = req.task_name
  90. with lock:
  91. status = cam_status.get(camera_id)
  92. if not status:
  93. return {"status": "not_running"}
  94. if status["task_name"] != task_name:
  95. return {"status": "task_mismatch", "running_task": status["task_name"]}
  96. # 标记停止
  97. status["stop_flag"] = True
  98. status["timer"].cancel()
  99. return {"status": "stopped", "file": status["file"]}
  100. @router.get("/status/{camera_id}")
  101. def status(camera_id: str):
  102. with lock:
  103. status = cam_status.get(camera_id)
  104. if not status:
  105. return {"status": "idle"}
  106. return {
  107. "status": "running",
  108. "task_name": status["task_name"],
  109. "start_time": status["start_time"],
  110. "file": status["file"]
  111. }