ROI选区01
This commit is contained in:
@@ -12,11 +12,11 @@ import numpy as np
|
||||
from config import get_config
|
||||
from db.crud import (
|
||||
create_alarm,
|
||||
get_all_rois,
|
||||
update_camera_status,
|
||||
)
|
||||
from db.models import init_db
|
||||
from db.models import init_db, get_session_factory
|
||||
from inference.engine import YOLOEngine
|
||||
from inference.roi.cache_manager import get_roi_cache
|
||||
from inference.roi.roi_filter import ROIFilter
|
||||
from inference.rules.algorithms import AlgorithmManager
|
||||
from inference.stream import StreamManager
|
||||
@@ -34,6 +34,8 @@ class InferencePipeline:
|
||||
reconnect_delay=self.config.stream.reconnect_delay,
|
||||
)
|
||||
self.roi_filter = ROIFilter()
|
||||
self.roi_cache = get_roi_cache()
|
||||
|
||||
self.algo_manager = AlgorithmManager(working_hours=[
|
||||
{
|
||||
"start": [wh.start[0], wh.start[1]],
|
||||
@@ -56,6 +58,8 @@ class InferencePipeline:
|
||||
if not self.db_initialized:
|
||||
init_db()
|
||||
self.db_initialized = True
|
||||
self.roi_cache.initialize(get_session_factory(), refresh_interval=10.0)
|
||||
self.roi_cache.start_background_refresh()
|
||||
|
||||
def _get_db_session(self):
|
||||
from db.models import get_session_factory
|
||||
@@ -169,85 +173,65 @@ class InferencePipeline:
|
||||
def _process_frame(self, camera_id: int, frame: np.ndarray, camera):
|
||||
from ultralytics.engine.results import Results
|
||||
|
||||
db = self._get_db_session()
|
||||
try:
|
||||
rois = get_all_rois(db, camera_id)
|
||||
roi_configs = [
|
||||
roi_configs = self.roi_cache.get_rois(camera_id)
|
||||
|
||||
if roi_configs:
|
||||
self.roi_filter.update_cache(camera_id, roi_configs)
|
||||
|
||||
for roi_config in roi_configs:
|
||||
roi_id = roi_config["roi_id"]
|
||||
rule_type = roi_config["rule"]
|
||||
|
||||
self.algo_manager.register_algorithm(
|
||||
roi_id,
|
||||
rule_type,
|
||||
{
|
||||
"id": roi.id,
|
||||
"roi_id": roi.roi_id,
|
||||
"type": roi.roi_type,
|
||||
"points": json.loads(roi.points),
|
||||
"rule": roi.rule_type,
|
||||
"direction": roi.direction,
|
||||
"enabled": roi.enabled,
|
||||
"threshold_sec": roi.threshold_sec,
|
||||
"confirm_sec": roi.confirm_sec,
|
||||
"return_sec": roi.return_sec,
|
||||
}
|
||||
for roi in rois
|
||||
]
|
||||
"threshold_sec": roi_config.get("threshold_sec", 360),
|
||||
"confirm_sec": roi_config.get("confirm_sec", 30),
|
||||
"return_sec": roi_config.get("return_sec", 5),
|
||||
},
|
||||
)
|
||||
|
||||
if roi_configs:
|
||||
self.roi_filter.update_cache(camera_id, roi_configs)
|
||||
results = self.yolo_engine(frame, verbose=False, classes=[0])
|
||||
|
||||
for roi_config in roi_configs:
|
||||
roi_id = roi_config["roi_id"]
|
||||
rule_type = roi_config["rule"]
|
||||
if not results:
|
||||
return
|
||||
|
||||
self.algo_manager.register_algorithm(
|
||||
result = results[0]
|
||||
detections = []
|
||||
if hasattr(result, "boxes") and result.boxes is not None:
|
||||
boxes = result.boxes.xyxy.cpu().numpy()
|
||||
confs = result.boxes.conf.cpu().numpy()
|
||||
for i, box in enumerate(boxes):
|
||||
detections.append({
|
||||
"bbox": box.tolist(),
|
||||
"conf": float(confs[i]),
|
||||
"cls": 0,
|
||||
})
|
||||
|
||||
if roi_configs:
|
||||
filtered_detections = self.roi_filter.filter_detections(
|
||||
detections, roi_configs
|
||||
)
|
||||
else:
|
||||
filtered_detections = detections
|
||||
|
||||
for detection in filtered_detections:
|
||||
matched_rois = detection.get("matched_rois", [])
|
||||
for roi_conf in matched_rois:
|
||||
roi_id = roi_conf["roi_id"]
|
||||
rule_type = roi_conf["rule"]
|
||||
|
||||
alerts = self.algo_manager.process(
|
||||
roi_id,
|
||||
str(camera_id),
|
||||
rule_type,
|
||||
{
|
||||
"threshold_sec": roi_config.get("threshold_sec", 360),
|
||||
"confirm_sec": roi_config.get("confirm_sec", 30),
|
||||
"return_sec": roi_config.get("return_sec", 5),
|
||||
},
|
||||
[detection],
|
||||
datetime.now(),
|
||||
)
|
||||
|
||||
results = self.yolo_engine(frame, verbose=False, classes=[0])
|
||||
|
||||
if not results:
|
||||
return
|
||||
|
||||
result = results[0]
|
||||
detections = []
|
||||
if hasattr(result, "boxes") and result.boxes is not None:
|
||||
boxes = result.boxes.xyxy.cpu().numpy()
|
||||
confs = result.boxes.conf.cpu().numpy()
|
||||
for i, box in enumerate(boxes):
|
||||
detections.append({
|
||||
"bbox": box.tolist(),
|
||||
"conf": float(confs[i]),
|
||||
"cls": 0,
|
||||
})
|
||||
|
||||
if roi_configs:
|
||||
filtered_detections = self.roi_filter.filter_detections(
|
||||
detections, roi_configs
|
||||
)
|
||||
else:
|
||||
filtered_detections = detections
|
||||
|
||||
for detection in filtered_detections:
|
||||
matched_rois = detection.get("matched_rois", [])
|
||||
for roi_conf in matched_rois:
|
||||
roi_id = roi_conf["roi_id"]
|
||||
rule_type = roi_conf["rule"]
|
||||
|
||||
alerts = self.algo_manager.process(
|
||||
roi_id,
|
||||
str(camera_id),
|
||||
rule_type,
|
||||
[detection],
|
||||
datetime.now(),
|
||||
)
|
||||
|
||||
for alert in alerts:
|
||||
self._handle_alert(camera_id, alert, frame, roi_conf)
|
||||
|
||||
finally:
|
||||
db.close()
|
||||
for alert in alerts:
|
||||
self._handle_alert(camera_id, alert, frame, roi_conf)
|
||||
|
||||
def _handle_alert(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user