weijuesen
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from urllib.parse import urlparse
|
|
|
|
from . import config
|
|
from .detector import MockDetector, ONNXDetector
|
|
|
|
|
|
def _detector():
|
|
if config.MODEL_MODE == "onnx":
|
|
return ONNXDetector(config.MODEL_PATH, tuple(config.MODEL_LABELS))
|
|
return MockDetector(class_name=config.MOCK_CLASS, confidence=config.MOCK_CONFIDENCE)
|
|
|
|
|
|
def is_allowed_stream_ref(stream_ref):
|
|
url = (stream_ref or {}).get("url", "")
|
|
parsed = urlparse(url)
|
|
if parsed.scheme not in ("rtsp", "http", "https"):
|
|
return False
|
|
if parsed.username or parsed.password:
|
|
return False
|
|
host = (parsed.hostname or "").lower()
|
|
allowed = {item.strip().lower() for item in config.ALLOWED_STREAM_HOSTS}
|
|
return host in allowed
|
|
|
|
|
|
class StreamTaskWorker:
|
|
def __init__(self, max_workers=2):
|
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
self.tasks = {}
|
|
self.lock = threading.Lock()
|
|
|
|
def submit(self, task_id, stream_ref):
|
|
with self.lock:
|
|
self.tasks[task_id] = {
|
|
"taskId": task_id,
|
|
"status": "queued",
|
|
"startedAt": time.time(),
|
|
"finishedAt": None,
|
|
"frames": 0,
|
|
"detections": [],
|
|
"error": None,
|
|
}
|
|
self.executor.submit(self._run, task_id, stream_ref)
|
|
|
|
def get(self, task_id):
|
|
with self.lock:
|
|
task = self.tasks.get(task_id)
|
|
return dict(task) if task else None
|
|
|
|
def _run(self, task_id, stream_ref):
|
|
try:
|
|
with self.lock:
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task["status"] = "running"
|
|
try:
|
|
import cv2
|
|
except ImportError:
|
|
self._fail(task_id, "OpenCV 未安装，无法拉流")
|
|
return
|
|
|
|
capture = cv2.VideoCapture(stream_ref.get("url", ""))
|
|
if not capture.isOpened():
|
|
self._fail(task_id, "无法连接视频流")
|
|
return
|
|
|
|
max_frames = min(int(stream_ref.get("maxFrames", config.STREAM_TASK_MAX_FRAMES)), config.STREAM_TASK_MAX_FRAMES)
|
|
detector = _detector()
|
|
frames = 0
|
|
detections = []
|
|
try:
|
|
while frames < max_frames:
|
|
ok, frame = capture.read()
|
|
if not ok:
|
|
break
|
|
ok_encode, buf = cv2.imencode(".jpg", frame)
|
|
if ok_encode:
|
|
detections.extend(detector.detect(buf.tobytes()))
|
|
frames += 1
|
|
finally:
|
|
capture.release()
|
|
|
|
with self.lock:
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task["status"] = "completed"
|
|
task["frames"] = frames
|
|
task["detections"] = detections[:10]
|
|
task["finishedAt"] = time.time()
|
|
except Exception as exc:
|
|
self._fail(task_id, str(exc))
|
|
|
|
def _fail(self, task_id, message):
|
|
with self.lock:
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task["status"] = "failed"
|
|
task["error"] = message
|
|
task["finishedAt"] = time.time()
|