THE githubnext FONT IS mad ugly bro

idfk how it could be that fucking chopped
This commit is contained in:
lda committed 2026-01-15 23:41:19 +07:00
1 parent d57a48eec6
commit cdbca0f370
7 files changed
+446 -277

No files matched your search

+2
View File
@@ -1,3 +1,5 @@
*.pt filter=lfs diff=lfs merge=lfs -text *.pt filter=lfs diff=lfs merge=lfs -text
*.mp4 filter=lfs diff=lfs merge=lfs -text *.mp4 filter=lfs diff=lfs merge=lfs -text
*.jsonl filter=lfs diff=lfs merge=lfs -text *.jsonl filter=lfs diff=lfs merge=lfs -text
*.png filter=lfs diff=lfs merge=lfs -text
*.jpg filter=lfs diff=lfs merge=lfs -text
+1
View File
@@ -6,4 +6,5 @@
"**/cache/": true "**/cache/": true
}, },
"pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"], "pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"],
// "editor.fontFamily": "'JetbrainsMono Nerd Font', Consolas, 'Courier New', monospace"
} }
+112 -34
View File
@@ -1,10 +1,10 @@
# Smoking Detection Pipeline # Smoking Detection Pipeline
A robust, multi-stage computer vision pipeline designed to detect smoking events in video files or real-time RTSP streams. The system uses a cascade of YOLO models to ensure high precision and low false positives. A robust, multi-stage computer vision pipeline designed to detect smoking events in video files, webcams, or real-time RTSP streams. The system uses a cascade of YOLO models with early-exit optimization for high precision and low latency.
## 🏗 Architecture ## 🏗 Architecture
The pipeline operates in a "Fail-Fast" cascade to optimize performance: The pipeline operates in a "Fail-Fast" cascade with **early exit** optimization:
1. **Broad Phase (Person Detection)**: 1. **Broad Phase (Person Detection)**:
@@ -13,13 +13,14 @@ The pipeline operates in a "Fail-Fast" cascade to optimize performance:
2. **Narrow Phase (Cigarette Candidate Detection)**: 2. **Narrow Phase (Cigarette Candidate Detection)**:
- Crops the detected person (with a 20% margin for context). - Crops the detected person (with configurable margin for context).
- Runs a specialized YOLO detector to find small objects (cigarettes/lighters). - Runs a specialized YOLO detector to find small objects (cigarettes/lighters).
- _Early Exit_: If confidence ≥ 0.8, skips verification phase entirely.
- _Filter_: Rejects candidates smaller than 10x10 pixels. - _Filter_: Rejects candidates smaller than 10x10 pixels.
3. **Verification Phase (Classification)**: 3. **Verification Phase (Classification)**:
- Crops the candidate object. - Crops the candidate object with padding.
- Runs a dedicated Binary Classifier to verify "Smoking" vs "Not Smoking". - Runs a dedicated Binary Classifier to verify "Smoking" vs "Not Smoking".
- _Output_: Only high-confidence events are flagged as alerts. - _Output_: Only high-confidence events are flagged as alerts.
@@ -27,12 +28,15 @@ The pipeline operates in a "Fail-Fast" cascade to optimize performance:
- **GPU Acceleration**: Auto-detects CUDA and runs inference on GPU when available. - **GPU Acceleration**: Auto-detects CUDA and runs inference on GPU when available.
- **Async Inference**: Detection runs in a background thread; video output never lags. - **Async Inference**: Detection runs in a background thread; video output never lags.
- **Early Exit**: High-confidence detections skip the verification stage for faster alerts.
- **Adaptive Sampling**: - **Adaptive Sampling**:
- **Idle Mode (2 FPS)**: When no people are detected, saves compute. - **Idle Mode (2 FPS)**: When no people are detected, saves compute.
- **Active Mode (Max FPS)**: Switches to full speed when activity is detected. - **Active Mode (Max FPS)**: Switches to full speed when activity is detected.
- **Cooldown**: Stays active for 2 seconds after people leave frame. - **Cooldown**: Stays active for 2 seconds after people leave frame.
- **Non-Maximum Suppression (NMS)**: Cleans up overlapping boxes per-class. - **Non-Maximum Suppression (NMS)**: Cleans up overlapping boxes per-class.
- **Smart Event Recording**: Captures 2s before + 3s after each smoking event as separate clips. - **Smart Event Recording**: Captures 2s before + 3s after each smoking event as separate clips.
- **Live Display**: Real-time visualization window with 'q' to quit.
- **Metrics Tracking**: FPS, latency, detection counts with summary on exit.
## 🚀 Usage ## 🚀 Usage
@@ -58,28 +62,42 @@ uv sync
# Run on a local video file # Run on a local video file
uv run main.py --source videos/test_footage.mp4 uv run main.py --source videos/test_footage.mp4
# Run on an RTSP stream # Run on webcam (index 0)
uv run main.py --source rtsp://admin:[email protected]:554/stream --rtsp uv run main.py --source 0 --display
# Enable debug logging (see frame-by-frame details) # Run on webcam with live display, no file output
uv run main.py --source 0 --display --no-save
# Run on an RTSP stream (auto-detected)
uv run main.py --source rtsp://admin:[email protected]:554/stream
# Enable debug logging
uv run main.py --source videos/test.mp4 --debug uv run main.py --source videos/test.mp4 --debug
``` ```
### Check Available Webcams
```bash
uv run -m tools.webcam_check
```
### CLI Arguments ### CLI Arguments
| Argument | Default | Description | | Argument | Default | Description |
| :--------------- | :------------------------- | :------------------------------------------------------------------------ | | :--------------- | :------------------------- | :-------------------------------------------------- |
| `--source` | **Required** | Path to video file or RTSP URL. | | `--source` | **Required** | Video file, webcam index (0, 1), or RTSP URL. |
| `--rtsp` | `False` | Set this flag if the source is an RTSP stream (enables threaded reading). | | `--rtsp` | `False` | Force RTSP mode (auto-detected for `rtsp://` URLs). |
| `--output` | `output.mp4` | Path for the annotated output video. | | `--output` | `output.mp4` | Path for the annotated output video. |
| `--person-model` | `models/human-det-best.pt` | Path to the Person Detection YOLO model. | | `--display` | `False` | Show live preview window (press 'q' to quit). |
| `--cig-model` | `models/cig-det-best.pt` | Path to the Cigarette Detection YOLO model. | | `--no-save` | `False` | Skip video/event output (useful with `--display`). |
| `--verify-model` | `models/cig-cls-best.pt` | Path to the Verification/Classifier model. | | `--person-model` | `models/human-det-best.pt` | Path to the Person Detection YOLO model. |
| `--debug` | `False` | Enable verbose logging. | | `--cig-model` | `models/cig-det-best.pt` | Path to the Cigarette Detection YOLO model. |
| `--verify-model` | `models/cig-cls-best.pt` | Path to the Verification/Classifier model. |
| `--debug` | `False` | Enable verbose logging. |
## 📊 Outputs ## 📊 Outputs
The system generates four types of artifacts during execution: The system generates multiple artifacts during execution:
1. **Visual Debug (`output.mp4`)**: 1. **Visual Debug (`output.mp4`)**:
@@ -91,12 +109,10 @@ The system generates four types of artifacts during execution:
2. **Frame Metadata (`output.jsonl`)**: 2. **Frame Metadata (`output.jsonl`)**:
- A synchronized JSON Lines file containing raw detection data for every frame. - A synchronized JSON Lines file containing raw detection data for every frame.
- Useful for debugging model performance or retraining.
3. **Alert Log (`alerts.jsonl`)**: 3. **Alert Log (`alerts.jsonl`)**:
- A structured log file containing **only** confirmed smoking events. - A structured log file containing **only** confirmed smoking events.
- Use this file to trigger external notifications or audit logic.
4. **Event Clips (`events/event_<timestamp>.mp4`)**: 4. **Event Clips (`events/event_<timestamp>.mp4`)**:
@@ -104,11 +120,20 @@ The system generates four types of artifacts during execution:
- Includes 2 seconds of pre-event context and 3 seconds post-event. - Includes 2 seconds of pre-event context and 3 seconds post-event.
- Automatically extends if smoking continues. - Automatically extends if smoking continues.
5. **Application Logs (`logs/pipeline.log`)**: 5. **Screenshots (`screenshots/alert_<frame>_<conf>.jpg`)**:
- Annotated JPG images saved when alerts fire.
- Useful for evidence/audit trails without full video.
6. **Application Logs (`logs/pipeline.log`)**:
- Rotating log files with detailed debug information. - Rotating log files with detailed debug information.
- Includes thread names, function calls, and timing data. - Includes thread names, function calls, and timing data.
7. **Metrics Summary** (printed on exit):
- Total frames, average FPS, latency, detection counts by label.
## 📂 Project Structure ## 📂 Project Structure
```text ```text
@@ -118,26 +143,79 @@ gp/
├── models/ # Directory for .pt model weights ├── models/ # Directory for .pt model weights
├── logs/ # Rotating application logs ├── logs/ # Rotating application logs
├── events/ # Smart event clips (auto-created) ├── events/ # Smart event clips (auto-created)
├── screenshots/ # Alert evidence images
│
├── pipeline/ ├── pipeline/
│ ├── core.py # Main logic loop (Pipeline class, Async Inference) │ ├── config.py # PipelineConfig dataclass (all thresholds)
│ ├── models.py # Ultralytics model wrappers (GPU-enabled) │ ├── core.py # SmokingDetectionPipeline (async inference)
│ ├── sources.py # Stream inputs (File, ThreadedRTSP)
│ ├── outputs.py # Sinks (VideoFileStore, SmartEventSink, FileAlertStore)
│ ├── protocols.py # Interfaces (Model, Verifier, Sink, StreamSource) │ ├── protocols.py # Interfaces (Model, Verifier, Sink, StreamSource)
│ └── types.py # Data classes (Frame, Detection, Ticket, NMS) │ ├── types.py # Data classes (Frame, Detection, Ticket, BoundingBox)
│ │
│ ├── models/ # ML model wrappers
│ │ └── yolo.py # YOLO person/cig/classifier (GPU-enabled)
│ │
│ ├── sources/ # Stream inputs
│ │ ├── file.py # FileStreamSource (video files, image sequences)
│ │ └── stream.py # ThreadedStreamSource (RTSP, webcam)
│ │
│ ├── sinks/ # Output handlers
│ │ ├── alerts.py # FileAlertStore (JSONL)
│ │ ├── display.py # WindowDisplaySink (live preview)
│ │ ├── events.py # SmartEventSink (clip recorder)
│ │ ├── metrics.py # MetricsSink (FPS, latency, counts)
│ │ ├── screenshots.py # ScreenshotSink (evidence images)
│ │ └── video.py # VideoFileStore (async video writer)
│ │
│ └── ops/ # Algorithms & utilities
│ ├── nms.py # Non-Maximum Suppression
│ ├── geometry.py # Box expansion, zone filtering
│ └── drawing.py # Visualization helpers
│
└── tools/
└── webcam_check.py # Utility to list available webcams
``` ```
## ⚙️ Performance Tuning ## ⚙️ Configuration
The pipeline is designed for real-time performance. Key parameters in `main.py`: All pipeline thresholds are centralized in `pipeline/config.py`:
```python ```python
pipeline = SmokingDetectionPipeline( @dataclass
... class PipelineConfig:
inference_interval=0.5, # Idle: 2 FPS (saves GPU when empty) # Detection thresholds
active_interval=0.0, # Active: Process every frame person_nms_iou: float = 0.7
cooldown_frames=60, # Stay active 2s after person leaves cig_nms_iou: float = 0.5
)
# Confidence thresholds
cig_min_confidence: float = 0.2 # Minimum to consider
cig_high_confidence: float = 0.8 # Skip verifier if above
verify_threshold: float = 0.8 # Verifier must exceed
# Adaptive sampling
idle_interval: float = 0.5 # 2 FPS when idle
active_interval: float = 0.0 # Max FPS when active
cooldown_frames: int = 60 # Stay active after detection
``` ```
For RTSP streams, the `ThreadedRTSPSource` uses a background thread to prevent frame buildup and handles reconnection automatically. Pass a custom config to the pipeline:
```python
from pipeline.config import PipelineConfig
config = PipelineConfig(
cig_high_confidence=0.9, # Stricter early exit
idle_interval=1.0, # 1 FPS when idle
)
pipeline = SmokingDetectionPipeline(..., config=config)
```
## 🔧 Performance Tuning
For RTSP streams, the `ThreadedStreamSource` uses a background thread to prevent frame buildup and handles reconnection automatically.
Key tuning parameters:
- `idle_interval`: Increase to save GPU when scene is empty
- `active_interval`: Set to 0.0 for max detection rate
- `cig_high_confidence`: Lower for faster alerts, higher for fewer false positives
+3 -2
View File
@@ -1,7 +1,8 @@
import cv2
import time
import threading import threading
import time
from collections import deque from collections import deque
import cv2
from ultralytics.models import YOLO from ultralytics.models import YOLO
# ========= CẤU HÌNH ========= # ========= CẤU HÌNH =========
+10 -6
View File
@@ -4,6 +4,7 @@ import logging.handlers
import sys import sys
from pathlib import Path from pathlib import Path
from pipeline.config import PipelineConfig
from pipeline.core import SmokingDetectionPipeline from pipeline.core import SmokingDetectionPipeline
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
from pipeline.sinks import ( from pipeline.sinks import (
@@ -165,17 +166,20 @@ def main():
display_sink = WindowDisplaySink("Smoking Detection", scale=1.0) display_sink = WindowDisplaySink("Smoking Detection", scale=1.0)
sinks.append(display_sink) sinks.append(display_sink)
# Build pipeline config
config = PipelineConfig(
idle_interval=0.5, # 2 FPS when nobody around
active_interval=0.0, # Max speed when person detected
cooldown_frames=int(source.fps * 2), # Stay active 2s after detection
)
pipeline = SmokingDetectionPipeline( pipeline = SmokingDetectionPipeline(
person_model=person_model, person_model=person_model,
cigarette_model=cigarette_model, cigarette_model=cigarette_model,
verifier_model=verifier, verifier_model=verifier,
vlm_model=None, # Optional vlm_model=None,
sinks=sinks, sinks=sinks,
inference_interval=0.5, # Slow mode (2 FPS) when nobody is home config=config,
active_interval=0.0, # Max speed when person detected
cooldown_frames=int(
source.fps * 2
), # Stay active for 2 seconds after person leaves
) )
logger.info("Starting pipeline with %d sinks...", len(sinks)) logger.info("Starting pipeline with %d sinks...", len(sinks))
+37
View File
@@ -0,0 +1,37 @@
"""Pipeline configuration."""
from dataclasses import dataclass
@dataclass
class PipelineConfig:
"""All pipeline thresholds and settings in one place."""
# Detection thresholds
person_nms_iou: float = 0.7
cig_nms_iou: float = 0.5
final_nms_iou: float = 0.5
# Confidence thresholds
cig_min_confidence: float = 0.2 # Minimum to consider a candidate
cig_high_confidence: float = 0.8 # Skip verifier if above this
verify_threshold: float = 0.8 # Verifier must exceed this
# Size filters (pixels)
min_crop_size: int = 10 # Minimum dimension for valid crop
min_cig_size: int = 10 # Minimum cig box dimension
# Person crop expansion
person_margin: float = 0.2 # 20% expansion around person box
# Verification crop padding
cig_crop_padding: float = 0.5 # 50% padding around cig for context
# Adaptive sampling
idle_interval: float = 0.5 # Seconds between inference when idle
active_interval: float = 0.0 # Seconds between inference when active
cooldown_frames: int = 60 # Frames to stay active after last detection
# Default config instance
DEFAULT_CONFIG = PipelineConfig()
+281 -235
View File
@@ -1,10 +1,14 @@
"""Smoking Detection Pipeline - Main orchestration."""
import concurrent.futures import concurrent.futures
import logging import logging
import time import time
from collections import deque from collections import deque
from concurrent.futures import Future from concurrent.futures import Future
from dataclasses import dataclass
from typing import List, Optional, Tuple from typing import List, Optional, Tuple
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
from pipeline.ops import nms_by_class from pipeline.ops import nms_by_class
from pipeline.protocols import Model, Sink, StreamSource, Verifier from pipeline.protocols import Model, Sink, StreamSource, Verifier
from pipeline.types import BoundingBox, Detection, Frame, Ticket from pipeline.types import BoundingBox, Detection, Frame, Ticket
@@ -12,7 +16,26 @@ from pipeline.types import BoundingBox, Detection, Frame, Ticket
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@dataclass
class InferenceResult:
"""Result of a single inference pass."""
has_activity: bool
detections: List[Detection]
ticket: Optional[Ticket] = None
class SmokingDetectionPipeline: class SmokingDetectionPipeline:
"""
Multi-stage smoking detection pipeline.
Stage 1: Person Detection (broad phase)
Stage 2: Cigarette Detection (narrow phase, per-person crops)
Stage 3: Verification (classifier confirmation)
Supports early exit at Stage 2 if cigarette confidence is high enough.
"""
def __init__( def __init__(
self, self,
person_model: Model, person_model: Model,
@@ -20,306 +43,329 @@ class SmokingDetectionPipeline:
verifier_model: Verifier, verifier_model: Verifier,
vlm_model: Optional[Verifier] = None, vlm_model: Optional[Verifier] = None,
sinks: Optional[list[Sink]] = None, sinks: Optional[list[Sink]] = None,
inference_interval: float = 0.5, config: Optional[PipelineConfig] = None,
active_interval: float = 0.0,
cooldown_frames: int = 60,
): ):
self.person_model = person_model self.person_model = person_model
self.cigarette_model = cigarette_model self.cigarette_model = cigarette_model
self.verifier_model = verifier_model self.verifier_model = verifier_model
self.vlm_model = vlm_model self.vlm_model = vlm_model
self.sinks = sinks or [] self.sinks = sinks or []
self.cfg = config or DEFAULT_CONFIG
# Adaptive Sampling Settings # Adaptive sampling state
self.idle_interval = inference_interval self.current_interval = self.cfg.idle_interval
self.active_interval = active_interval
self.cooldown_frames = cooldown_frames
# State
self.current_interval = self.idle_interval
self.active_cooldown = 0 self.active_cooldown = 0
self.last_inference_time = 0.0 self.last_inference_time = 0.0
# Async Inference State # Async inference state
self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
self.current_future: Optional[Future] = None self.current_future: Optional[Future[InferenceResult]] = None
self.latest_detections: List[ self.latest_result: Optional[InferenceResult] = None
Detection
] = [] # Cache for display to prevent flickering
self.latest_ticket: Optional[Ticket] = None
# Buffer for temporal context (optional, for VLM or sequence logic) # Frame buffer for temporal context
self.frame_buffer = deque(maxlen=30) self.frame_buffer = deque(maxlen=30)
logger.info( logger.info(
f"Pipeline initialized with {len(self.sinks)} sinks. Async Inference enabled." "Pipeline initialized with %d sinks. Config: %s",
len(self.sinks),
self.cfg,
) )
def process_stream(self, source: StreamSource): def process_stream(self, source: StreamSource):
"""Main loop processing the video stream.""" """Main loop: read frames, schedule inference, dispatch to sinks."""
logger.info("Starting stream processing...") logger.info("Starting stream processing...")
count = 0 frame_count = 0
for frame in source.stream(): for frame in source.stream():
count += 1 frame_count += 1
if count % 100 == 0: if frame_count % 100 == 0:
logger.info(f"Processed {count} frames...") logger.info("Processed %d frames...", frame_count)
self.frame_buffer.append(frame) self.frame_buffer.append(frame)
self._poll_inference_result(frame)
self._maybe_schedule_inference(frame)
self._apply_detections_to_frame(frame)
self._dispatch_to_sinks(frame)
# 1. Check if previous inference finished logger.info("Stream finished. Total frames: %d", frame_count)
if self.current_future and self.current_future.done():
try:
# Retrieve results from the finished job
has_activity, detections, ticket = self.current_future.result()
# Update Cache
self.latest_detections = detections
self.latest_ticket = ticket
self.current_future = None
# Adaptive Logic Update
if has_activity:
if self.current_interval != self.active_interval:
logger.info(
"Frame %s: Activity detected! Switching to ACTIVE mode (0s interval).",
frame.frame_id,
)
self.current_interval = self.active_interval
self.active_cooldown = self.cooldown_frames
elif self.active_cooldown > 0:
if self.active_cooldown % 30 == 0:
logger.debug(
"Cooling down... %s frames left.", self.active_cooldown
)
self.active_cooldown -= 1
else:
if self.current_interval != self.idle_interval:
logger.info(
"Cooldown complete. Switching to IDLE mode (0.5s interval)."
)
self.current_interval = self.idle_interval
except Exception as e:
logger.error(f"Inference error: {e}")
self.current_future = None
# 2. Schedule new inference if idle and time condition met
now = frame.timestamp
time_since_last = now - self.last_inference_time
if self.current_future is None and time_since_last >= self.current_interval:
# Submit 'detect' job. We pass a COPY of data or rely on immutable usage.
# Frame.data is numpy, which is mutable but we trust models not to write to it.
# Since models might run on GPU/Thread, and we continue, it's safer to not rely on 'frame' object mutation.
# We will return the results and apply them later.
self.last_inference_time = now
self.current_future = self.executor.submit(self._run_inference, frame)
# 3. Apply cached detections to current frame for visualization
# This ensures smooth video output even if inference runs at lower fps
frame.detections = self.latest_detections.copy()
# Ticket dispatch: Only attach to ONE frame, then clear
# This prevents duplicate alerts in sinks like FileAlertStore
if self.latest_ticket:
frame.ticket = self.latest_ticket
self.latest_ticket = None # Consume the ticket
# Update sinks (Visualizers, etc)
for sink in self.sinks:
sink.handle(frame)
logger.info("Stream processing finished.")
self.executor.shutdown(wait=False) self.executor.shutdown(wait=False)
def _run_inference( def _poll_inference_result(self, frame: Frame):
self, frame: Frame """Check if async inference completed, update state."""
) -> Tuple[bool, List[Detection], Optional[Ticket]]: if not self.current_future or not self.current_future.done():
return
try:
self.latest_result = self.current_future.result()
self._update_adaptive_sampling(frame, self.latest_result.has_activity)
except Exception as e:
logger.error("Inference error: %s", e)
finally:
self.current_future = None
def _update_adaptive_sampling(self, frame: Frame, has_activity: bool):
"""Switch between idle/active modes based on detections."""
if has_activity:
if self.current_interval != self.cfg.active_interval:
logger.info(
"Frame %d: Activity! Switching to ACTIVE mode.", frame.frame_id
)
self.current_interval = self.cfg.active_interval
self.active_cooldown = self.cfg.cooldown_frames
elif self.active_cooldown > 0:
self.active_cooldown -= 1
else:
if self.current_interval != self.cfg.idle_interval:
logger.info("Cooldown complete. Switching to IDLE mode.")
self.current_interval = self.cfg.idle_interval
def _maybe_schedule_inference(self, frame: Frame):
"""Schedule inference if enough time has passed."""
if self.current_future is not None:
return
time_since_last = frame.timestamp - self.last_inference_time
if time_since_last >= self.current_interval:
self.last_inference_time = frame.timestamp
self.current_future = self.executor.submit(self._run_inference, frame)
def _apply_detections_to_frame(self, frame: Frame):
"""Apply cached detections to current frame for visualization."""
if self.latest_result:
frame.detections = self.latest_result.detections.copy()
# Consume ticket (only dispatch once)
if self.latest_result.ticket:
frame.ticket = self.latest_result.ticket
self.latest_result.ticket = None
def _dispatch_to_sinks(self, frame: Frame):
"""Send frame to all registered sinks."""
for sink in self.sinks:
sink.handle(frame)
# -------------------------------------------------------------------------
# Inference Chain
# -------------------------------------------------------------------------
def _run_inference(self, frame: Frame) -> InferenceResult:
""" """
Runs the full detection chain. Full detection chain. Thread-safe (doesn't mutate frame).
Returns (ActivityDetected, ListOfDetections, TicketOrNone).
Does NOT modify frame in place (thread safety). Returns InferenceResult with detections and optional ticket.
""" """
t0 = time.time() t0 = time.time()
detections = [] detections: List[Detection] = []
# --- Step 1: Broad Phase (Person Detection) ---
people = self.person_model.predict(frame.data)
people = nms_by_class(people, iou_threshold=0.7)
# Stage 1: Person Detection
people = self._detect_people(frame)
detections.extend(people) detections.extend(people)
if not people: if not people:
# logger.debug("No people found.") return InferenceResult(has_activity=False, detections=detections)
return False, detections, None
# --- Step 2: Narrow Phase (Cigarette Candidate) --- # Stage 2: Cigarette Detection (with early exit for high confidence)
candidates = self._find_candidates(frame, people, detections) candidates, early_hit = self._detect_cigarettes(frame, people, detections)
if early_hit:
# High-confidence cig detection, skip verifier
ticket = self._create_ticket(frame, early_hit, stage="Detector")
detections = nms_by_class(detections, self.cfg.final_nms_iou)
self._log_inference_time(t0)
return InferenceResult(
has_activity=True, detections=detections, ticket=ticket
)
if not candidates: if not candidates:
return True, detections, None return InferenceResult(has_activity=True, detections=detections)
# --- Step 3: Verification (Classifier) --- # Stage 3: Verification
best_hit = self._verify_candidates(frame, candidates, detections) verified_hit = self._verify_candidates(frame, candidates, detections)
ticket = None ticket = None
# --- Step 4: Action --- if verified_hit:
if best_hit: ticket = self._create_ticket(frame, verified_hit, stage="Verifier")
person, cig, score = best_hit
logger.info("CONFIRMED SMOKING! Score: %.2f", score)
ticket = Ticket(
frame_id=frame.frame_id,
confidence=score,
evidence=[frame.data], # Note: this keeps ref to image data
stage_confirmed="Verifier",
)
# Add the verified box to detections for visualization detections = nms_by_class(detections, self.cfg.final_nms_iou)
verified_det = Detection( self._log_inference_time(t0)
box=BoundingBox( return InferenceResult(has_activity=True, detections=detections, ticket=ticket)
x1=person.box.x1
+ cig.box.x1, # Approx logic, need exact coords from find_candidates
y1=person.box.y1 + cig.box.y1,
x2=person.box.x1 + cig.box.x2,
y2=person.box.y1 + cig.box.y2,
conf=score, # Placeholder
),
score=score,
label="verified_smoking",
)
detections.append(verified_det)
# Clean up overlapping boxes from model output def _detect_people(self, frame: Frame) -> List[Detection]:
detections = nms_by_class(detections, iou_threshold=0.5) """Stage 1: Detect people in frame."""
people = self.person_model.predict(frame.data)
return nms_by_class(people, self.cfg.person_nms_iou)
dt = time.time() - t0 def _detect_cigarettes(
if dt > 0.1: self,
logger.debug("Inference took %.3fs (Slow)", dt) frame: Frame,
people: List[Detection],
detections_out: List[Detection],
) -> Tuple[List[Tuple], Optional[Tuple]]:
"""
Stage 2: Detect cigarettes in person crops.
return True, detections, ticket Returns:
(candidates, early_hit)
def _find_candidates( - candidates: list of (person, cigs, offset) for verification
self, frame: Frame, people: List[Detection], detections_out: List[Detection] - early_hit: (person, cig, score, offset) if high-confidence detection, else None
) -> List[Tuple[Detection, List[Detection], Tuple[int, int]]]:
"""Runs cigarette detection on expanded person crops.
Args:
detections_out: List to append candidate detections to (thread-safe output)
""" """
candidates = [] candidates = []
early_hit = None
h_img, w_img = frame.data.shape[:2] h_img, w_img = frame.data.shape[:2]
margin = 0.2 # 20% expansion cfg = self.cfg
for person in people: for person in people:
# 1. Expand Box # Expand person box
pb = person.box crop_box = self._expand_box(person.box, cfg.person_margin, w_img, h_img)
bw = pb.x2 - pb.x1 ex1, ey1, ex2, ey2 = crop_box
bh = pb.y2 - pb.y1
ex1 = max(0, int(pb.x1 - bw * margin))
ey1 = max(0, int(pb.y1 - bh * margin))
ex2 = min(w_img, int(pb.x2 + bw * margin))
ey2 = min(h_img, int(pb.y2 + bh * margin))
# 2. Crop
person_crop = frame.data[ey1:ey2, ex1:ex2] person_crop = frame.data[ey1:ey2, ex1:ex2]
if ( if not self._is_valid_crop(person_crop):
person_crop.size == 0
or person_crop.shape[0] < 10
or person_crop.shape[1] < 10
):
continue continue
# 3. Detect # Detect cigarettes in crop
cig_candidates = self.cigarette_model.predict(person_crop) cigs = self.cigarette_model.predict(person_crop)
cigs = self._filter_small_boxes(cigs, cfg.min_cig_size)
cigs = nms_by_class(cigs, cfg.cig_nms_iou)
# 4. Filter Noise + NMS (model can return overlapping boxes for same cig) for cig in cigs:
cig_candidates = [ # Add to visualization
c global_det = self._to_global_coords(cig, ex1, ey1, "candidate_cig")
for c in cig_candidates detections_out.append(global_det)
if (c.box.x2 - c.box.x1) >= 10 and (c.box.y2 - c.box.y1) >= 10
]
cig_candidates = nms_by_class(cig_candidates, iou_threshold=0.5)
if cig_candidates: # Check for early exit (high confidence)
candidates.append((person, cig_candidates, (ex1, ey1))) if cig.score >= cfg.cig_high_confidence and early_hit is None:
early_hit = (person, cig, cig.score, (ex1, ey1))
# 5. Global Coordinate Mapping (For Visualization) logger.info(
for cig in cig_candidates: "High-confidence cig detection (%.2f) - skipping verifier",
global_cig = Detection( cig.score,
box=BoundingBox(
x1=ex1 + cig.box.x1,
y1=ey1 + cig.box.y1,
x2=ex1 + cig.box.x2,
y2=ey1 + cig.box.y2,
conf=cig.box.conf,
),
score=cig.score,
label="candidate_cig",
metadata=cig.metadata,
) )
detections_out.append(global_cig)
return candidates if cigs:
candidates.append((person, cigs, (ex1, ey1)))
return candidates, early_hit
def _verify_candidates( def _verify_candidates(
self, frame: Frame, candidates: List[Tuple], detections_out: List[Detection] self,
frame: Frame,
candidates: List[Tuple],
detections_out: List[Detection],
) -> Optional[Tuple]: ) -> Optional[Tuple]:
"""Runs verification on candidate crops.
Args:
detections_out: List to append verified detections to (thread-safe output)
""" """
confirmed_hits = [] Stage 3: Run classifier on cigarette crops.
Returns best (person, cig, score, offset) if verified, else None.
"""
confirmed = []
h_img, w_img = frame.data.shape[:2]
cfg = self.cfg
for person, cigs, (ex1, ey1) in candidates: for person, cigs, (ex1, ey1) in candidates:
for cig in cigs: for cig in cigs:
# 1. Derive Crop Coordinates # Already caught by early exit?
cig_w = cig.box.x2 - cig.box.x1 if cig.score >= cfg.cig_high_confidence:
cig_h = cig.box.y2 - cig.box.y1
pad_x = int(cig_w * 0.5)
pad_y = int(cig_h * 0.5)
global_cx1 = ex1 + max(0, cig.box.x1 - pad_x)
global_cy1 = ey1 + max(0, cig.box.y1 - pad_y)
global_cx2 = ex1 + cig.box.x2 + pad_x
global_cy2 = ey1 + cig.box.y2 + pad_y
# Check bounds
h_img, w_img = frame.data.shape[:2]
final_x1 = max(0, global_cx1)
final_y1 = max(0, global_cy1)
final_x2 = min(w_img, global_cx2)
final_y2 = min(h_img, global_cy2)
# 2. Crop
obj_crop = frame.data[final_y1:final_y2, final_x1:final_x2]
if obj_crop.size == 0:
continue continue
# 3. Verify # Extract padded crop for classifier
score = self.verifier_model.verify(obj_crop) crop = self._extract_cig_crop(frame.data, cig, ex1, ey1, w_img, h_img)
if crop.size == 0:
continue
if score > 0.8: # High confidence score = self.verifier_model.verify(crop)
confirmed_hits.append((person, cig, score))
verified_det = Detection( if score >= cfg.verify_threshold:
box=BoundingBox( confirmed.append((person, cig, score, (ex1, ey1)))
x1=ex1 + cig.box.x1,
y1=ey1 + cig.box.y1, # Add verified detection
x2=ex1 + cig.box.x2, verified_det = self._to_global_coords(
y2=ey1 + cig.box.y2, cig, ex1, ey1, "verified_smoking"
conf=score,
),
score=score,
label="verified_smoking",
metadata={"verifier_score": score},
) )
verified_det.score = score
verified_det.metadata["verifier_score"] = score
detections_out.append(verified_det) detections_out.append(verified_det)
if confirmed_hits: return max(confirmed, key=lambda x: x[2]) if confirmed else None
return max(confirmed_hits, key=lambda x: x[2])
return None # -------------------------------------------------------------------------
# Helpers
# -------------------------------------------------------------------------
def _expand_box(
self, box: BoundingBox, margin: float, w_max: int, h_max: int
) -> Tuple[int, int, int, int]:
"""Expand box by margin, clamped to image bounds."""
bw, bh = box.x2 - box.x1, box.y2 - box.y1
x1 = max(0, int(box.x1 - bw * margin))
y1 = max(0, int(box.y1 - bh * margin))
x2 = min(w_max, int(box.x2 + bw * margin))
y2 = min(h_max, int(box.y2 + bh * margin))
return x1, y1, x2, y2
def _is_valid_crop(self, crop) -> bool:
"""Check if crop is large enough to process."""
return (
crop.size > 0
and crop.shape[0] >= self.cfg.min_crop_size
and crop.shape[1] >= self.cfg.min_crop_size
)
def _filter_small_boxes(
self, detections: List[Detection], min_size: int
) -> List[Detection]:
"""Remove boxes smaller than min_size."""
return [
d
for d in detections
if (d.box.x2 - d.box.x1) >= min_size and (d.box.y2 - d.box.y1) >= min_size
]
def _to_global_coords(
self, det: Detection, offset_x: int, offset_y: int, label: str
) -> Detection:
"""Convert detection from crop coords to global frame coords."""
return Detection(
box=BoundingBox(
x1=offset_x + det.box.x1,
y1=offset_y + det.box.y1,
x2=offset_x + det.box.x2,
y2=offset_y + det.box.y2,
conf=det.box.conf,
),
score=det.score,
label=label,
metadata=det.metadata.copy(),
)
def _extract_cig_crop(
self, image, cig: Detection, ex1: int, ey1: int, w_max: int, h_max: int
):
"""Extract padded crop around cigarette for verification."""
pad = self.cfg.cig_crop_padding
cig_w = cig.box.x2 - cig.box.x1
cig_h = cig.box.y2 - cig.box.y1
# Global coords with padding
x1 = max(0, ex1 + cig.box.x1 - int(cig_w * pad))
y1 = max(0, ey1 + cig.box.y1 - int(cig_h * pad))
x2 = min(w_max, ex1 + cig.box.x2 + int(cig_w * pad))
y2 = min(h_max, ey1 + cig.box.y2 + int(cig_h * pad))
return image[y1:y2, x1:x2]
def _create_ticket(self, frame: Frame, hit: Tuple, stage: str) -> Ticket:
"""Create alert ticket from detection hit."""
person, cig, score, (ex1, ey1) = hit
logger.info("CONFIRMED SMOKING! Score: %.2f (Stage: %s)", score, stage)
return Ticket(
frame_id=frame.frame_id,
confidence=score,
evidence=[frame.data],
stage_confirmed=stage,
)
def _log_inference_time(self, t0: float):
"""Log if inference was slow."""
dt = time.time() - t0
if dt > 0.1:
logger.debug("Inference took %.3fs (slow)", dt)