diff --git a/README.md b/README.md index 41d396b6..551a2de2 100644 --- a/README.md +++ b/README.md @@ -128,6 +128,10 @@ A `./data` directory is expected with at least: * optionally `model.onnx` to override weights from Hugging Face * optionally `config.json` to override model configuration +For cameras that repeat their latest snapshot with a fixed model, `--reuse-identical-frames` reuses detections when the +uncompressed pixels and model/mask settings match. Each pass still updates alert history and uploads. +This option is off by default; fresh images still run the model. + --- ## Camera configuration and adapters diff --git a/pyroengine/engine.py b/pyroengine/engine.py index 3219f3db..ac2b8f58 100644 --- a/pyroengine/engine.py +++ b/pyroengine/engine.py @@ -3,6 +3,7 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. +import hashlib import io import logging import shutil @@ -102,6 +103,7 @@ class Engine(Predictor): alert_relaxation: number of consecutive positive detections required to send the first alert, and also the number of consecutive negative detections before stopping the alert frame_size: Resize frame to frame_size before sending it to the api in order to save bandwidth (H, W) + reuse_identical_frames: reuse detections for unchanged camera pixels and model/mask settings cache_backup_period: number of minutes between each cache backup to disk frame_saving_period: Send one frame over N to the api for our dataset cache_size: maximum number of alerts to save in cache @@ -140,6 +142,7 @@ def __init__( save_detections_frames: Optional[bool] = False, send_last_image_period: int = 3600, # 1H last_bbox_mask_fetch_period: int = 3600, # 1H + reuse_identical_frames: bool = False, **kwargs: Any, # noqa: ANN401 ) -> None: cam_ids = list(cam_creds.keys()) if isinstance(cam_creds, dict) else None @@ -165,6 +168,7 @@ def __init__( # Cache & relaxation self.frame_saving_period = frame_saving_period + self.reuse_identical_frames = reuse_identical_frames self.jpeg_quality = jpeg_quality self.cache_backup_period = cache_backup_period self.day_time_strategy = day_time_strategy @@ -318,11 +322,29 @@ def predict( except RequestException as e: logger.warning(f"Failed to fetch occlusion masks for cam {cam_key} (pose {pose_id}): {e}") - # Inference with ONNX + inference_key = None if fake_pred is None: bbox_mask_dict = self.occlusion_masks.get(cam_key, {}) rgb_frame = frame if frame.mode == "RGB" else frame.convert("RGB") - preds = self.model(rgb_frame, bbox_mask_dict) + if self.reuse_identical_frames: + # JPEG equality can hide changes in the uncompressed pixels fed to the model. + inference_key = ( + rgb_frame.size, + hashlib.sha256(rgb_frame.tobytes()).digest(), + id(self.model), + id(self.model.model if self.model.format == "ncnn" else self.model.ort_session), + self.model.format, + repr(self.model.imgsz), + self.model.conf, + self.model.iou, + self.model.max_bbox_size, + tuple(tuple(box[:4]) for box in bbox_mask_dict.values()), + ) + history = self._states[cam_key]["last_predictions"] + if inference_key is not None and history and inference_key == self._states[cam_key].get("inference_key"): + preds = history[-1][1].copy() + else: + preds = self.model(rgb_frame, bbox_mask_dict) else: if fake_pred.size == 0: preds = np.empty((0, 5)) @@ -342,6 +364,7 @@ def predict( extra_boxes = state["event_crop_boxes"] if state["ongoing"] else None context_crop = self._build_context_crop(original_frame, preds, extra_boxes) conf = self._update_states(context_crop, preds, cam_key, encoded_bytes=encoded_bytes) + state["inference_key"] = inference_key if not self._states[cam_key]["ongoing"]: self._end_event(cam_key) diff --git a/src/run.py b/src/run.py index 55cabdf4..75691889 100644 --- a/src/run.py +++ b/src/run.py @@ -75,6 +75,7 @@ def main(args): backup_size=args.backup_size, nb_consecutive_frames=args.nb_consecutive_frames, frame_size=args.frame_size, + reuse_identical_frames=args.reuse_identical_frames, cache_backup_period=args.cache_backup_period, cache_size=args.cache_size, jpeg_quality=args.jpeg_quality, @@ -95,6 +96,9 @@ def main(args): ) # Model parser.add_argument("--model_path", type=str, default=None, help="model path") + parser.add_argument( + "--reuse-identical-frames", action="store_true", help="Reuse detections for identical camera snapshots" + ) parser.add_argument("--thresh", type=float, default=0.25, help="Confidence threshold") parser.add_argument("--max_bbox_size", type=float, default=0.4, help="Maximum bbox size") diff --git a/tests/test_engine_reuse.py b/tests/test_engine_reuse.py new file mode 100644 index 00000000..dd0f1602 --- /dev/null +++ b/tests/test_engine_reuse.py @@ -0,0 +1,96 @@ +import io +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from PIL import Image + +from pyroengine.engine import Engine + + +@pytest.fixture +def engine_factory(monkeypatch, tmp_path): + def classifier(**kwargs): + model = MagicMock(format="ncnn", imgsz=1024, conf=kwargs["conf"], iou=0, max_bbox_size=0.4) + model.model = object() + + def predict(frame, masks): + score = frame.getpixel((0, 0))[0] / 255 + return ( + np.array([[0.1, 0.2, 0.3, 0.4, score]], dtype=np.float32) + if score > model.conf and not masks + else np.empty((0, 5), dtype=np.float32) + ) + + model.side_effect = predict + return model + + monkeypatch.setattr("pyro_predictor.predictor.Classifier", classifier) + return lambda reuse=False: Engine( + cache_folder=str(tmp_path), nb_consecutive_frames=4, reuse_identical_frames=reuse, verbose=False + ) + + +def test_reuse_preserves_per_camera_alert_history(engine_factory): + baseline, cached = engine_factory(), engine_factory(True) + fire, forest = Image.new("RGB", (32, 32), "red"), Image.new("RGB", (32, 32)) + for cam, frame in [("a", fire), ("b", forest)] * 4 + [("a", forest)] * 4: + assert baseline.predict(frame.copy(), cam) == cached.predict(frame.copy(), cam) + a, b = baseline._states[cam], cached._states[cam] + assert a["ongoing"] == b["ongoing"] + assert len(a["last_predictions"]) == len(b["last_predictions"]) + for left, right in zip(a["last_predictions"], b["last_predictions"]): + np.testing.assert_array_equal(left[1], right[1]) + assert left[2] == right[2] + assert left[5] == right[5] + if len(b["last_predictions"]) > 1: + assert not np.shares_memory(b["last_predictions"][-1][1], b["last_predictions"][-2][1]) + assert baseline.model.call_count == 12 + assert cached.model.call_count == 3 + + +@pytest.mark.parametrize("change", ["pixels", "size", "mask", "threshold", "backend"]) +def test_reuse_invalidates_changed_inputs(engine_factory, change): + engine = engine_factory(True) + frame = Image.new("RGB", (32, 32)) + masks = {"area": [0.0, 0.0, 0.1, 0.1]} + engine.occlusion_masks["a"] = masks + engine.predict(frame, "a") + engine.predict(frame.copy(), "a") + assert engine.model.call_count == 1 + if change == "pixels": + before, after = io.BytesIO(), io.BytesIO() + frame.save(before, "JPEG", quality=80) + frame.putpixel((0, 0), (1, 0, 0)) + frame.save(after, "JPEG", quality=80) + assert before.getvalue() == after.getvalue() + elif change == "size": + frame = Image.new("RGB", (64, 16)) + elif change == "mask": + masks["area"][0] = 0.05 + elif change == "threshold": + engine.model.conf = 0.9 + else: + engine.model.model = object() + engine.predict(frame, "a") + assert engine.model.call_count == 2 + + +def test_fake_predictions_do_not_seed_reuse(engine_factory): + engine = engine_factory(True) + frame = Image.new("RGB", (32, 32), "red") + engine.predict(frame) + engine.predict(frame, fake_pred=np.empty((0, 5))) + engine.predict(frame) + assert engine.model.call_count == 2 + + +def test_failed_prediction_does_not_seed_reuse(engine_factory): + engine = engine_factory(True) + engine.predict(Image.new("RGB", (32, 32))) + frame = Image.new("RGB", (32, 32), "red") + with patch.object(engine, "_build_context_crop", side_effect=RuntimeError("crop failed")): + with pytest.raises(RuntimeError, match="crop failed"): + engine.predict(frame) + engine.predict(frame) + assert engine.model.call_count == 3