jexck/gp
1
0 Fork 0
forked from usth/gp
Code Pull Requests Activity

THE githubnext FONT IS mad ugly bro

Browse Source 
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
+447 -278
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitattributes
+2
View File
@@ -1,3 +1,5 @@
*.pt filter=lfs diff=lfs merge=lfs -text
*.mp4 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
.vscode/settings.json
Vendored
+1
View File
@@ -6,4 +6,5 @@
"**/cache/": true
},
"pylint.args": ["--generated-members=cv2.*", "--extension-pkg-whitelist=cv2"],
// "editor.fontFamily": "'JetbrainsMono Nerd Font', Consolas, 'Courier New', monospace"
}
README.md
+113 -35
View File
@@ -1,10 +1,10 @@
# 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
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)**:
@@ -13,13 +13,14 @@ The pipeline operates in a "Fail-Fast" cascade to optimize performance:
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).
- _Early Exit_: If confidence ≥ 0.8, skips verification phase entirely.
- _Filter_: Rejects candidates smaller than 10x10 pixels.
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".
- _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.
- **Async Inference**: Detection runs in a background thread; video output never lags.
- **Adaptive Sampling**:
- **Early Exit**: High-confidence detections skip the verification stage for faster alerts.
- **Adaptive Sampling**:
- **Idle Mode (2 FPS)**: When no people are detected, saves compute.
- **Active Mode (Max FPS)**: Switches to full speed when activity is detected.
- **Cooldown**: Stays active for 2 seconds after people leave frame.
- **Non-Maximum Suppression (NMS)**: Cleans up overlapping boxes per-class.
- **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
@@ -58,28 +62,42 @@ uv sync
# Run on a local video file
uv run main.py --source videos/test_footage.mp4
# Run on an RTSP stream
uv run main.py --source rtsp://admin:[email protected]:554/stream --rtsp
# Run on webcam (index 0)
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
```
### Check Available Webcams
```bash
uv run -m tools.webcam_check
```
### CLI Arguments
| Argument | Default | Description |
| :--------------- | :------------------------- | :------------------------------------------------------------------------ |
| `--source` | **Required** | Path to video file or RTSP URL. |
| `--rtsp` | `False` | Set this flag if the source is an RTSP stream (enables threaded reading). |
| `--output` | `output.mp4` | Path for the annotated output video. |
| `--person-model` | `models/human-det-best.pt` | Path to the Person Detection YOLO model. |
| `--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. |
| Argument | Default | Description |
| :--------------- | :------------------------- | :-------------------------------------------------- |
| `--source` | **Required** | Video file, webcam index (0, 1), or RTSP URL. |
| `--rtsp` | `False` | Force RTSP mode (auto-detected for `rtsp://` URLs). |
| `--output` | `output.mp4` | Path for the annotated output video. |
| `--display` | `False` | Show live preview window (press 'q' to quit). |
| `--no-save` | `False` | Skip video/event output (useful with `--display`). |
| `--person-model` | `models/human-det-best.pt` | Path to the Person Detection YOLO model. |
| `--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
The system generates four types of artifacts during execution:
The system generates multiple artifacts during execution:
1. **Visual Debug (`output.mp4`)**:
@@ -91,12 +109,10 @@ The system generates four types of artifacts during execution:
2. **Frame Metadata (`output.jsonl`)**:
- A synchronized JSON Lines file containing raw detection data for every frame.
- Useful for debugging model performance or retraining.
3. **Alert Log (`alerts.jsonl`)**:
- 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`)**:
@@ -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.
- 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.
- 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
```text
@@ -118,26 +143,79 @@ gp/
├── models/ # Directory for .pt model weights
├── logs/ # Rotating application logs
├── events/ # Smart event clips (auto-created)
├── screenshots/ # Alert evidence images
│
├── pipeline/
│ ├── core.py # Main logic loop (Pipeline class, Async Inference)
│ ├── models.py # Ultralytics model wrappers (GPU-enabled)
│ ├── sources.py # Stream inputs (File, ThreadedRTSP)
│ ├── outputs.py # Sinks (VideoFileStore, SmartEventSink, FileAlertStore)
│ ├── config.py # PipelineConfig dataclass (all thresholds)
│ ├── core.py # SmokingDetectionPipeline (async inference)
│ ├── 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
pipeline = SmokingDetectionPipeline(
...
inference_interval=0.5, # Idle: 2 FPS (saves GPU when empty)
active_interval=0.0, # Active: Process every frame
cooldown_frames=60, # Stay active 2s after person leaves
)
@dataclass
class PipelineConfig:
# Detection thresholds
person_nms_iou: float = 0.7
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
legacy/test_cam.py
+3 -2
View File
@@ -1,7 +1,8 @@
import cv2
import time
import threading
import time
from collections import deque
import cv2
from ultralytics.models import YOLO
# ========= CẤU HÌNH =========
main.py
+10 -6
View File
@@ -4,6 +4,7 @@ import logging.handlers
import sys
from pathlib import Path
from pipeline.config import PipelineConfig
from pipeline.core import SmokingDetectionPipeline
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
from pipeline.sinks import (
@@ -165,17 +166,20 @@ def main():
display_sink = WindowDisplaySink("Smoking Detection", scale=1.0)
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(
person_model=person_model,
cigarette_model=cigarette_model,
verifier_model=verifier,
vlm_model=None, # Optional
vlm_model=None,
sinks=sinks,
inference_interval=0.5, # Slow mode (2 FPS) when nobody is home
active_interval=0.0, # Max speed when person detected
cooldown_frames=int(
source.fps * 2
), # Stay active for 2 seconds after person leaves
config=config,
)
logger.info("Starting pipeline with %d sinks...", len(sinks))
pipeline/config.py
+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()
pipeline/core.py
+281 -235
View File
@@ -1,10 +1,14 @@
"""Smoking Detection Pipeline - Main orchestration."""
import concurrent.futures
import logging
import time
from collections import deque
from concurrent.futures import Future
from dataclasses import dataclass
from typing import List, Optional, Tuple
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
from pipeline.ops import nms_by_class
from pipeline.protocols import Model, Sink, StreamSource, Verifier
from pipeline.types import BoundingBox, Detection, Frame, Ticket
@@ -12,7 +16,26 @@ from pipeline.types import BoundingBox, Detection, Frame, Ticket
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:
"""
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__(
self,
person_model: Model,
@@ -20,306 +43,329 @@ class SmokingDetectionPipeline:
verifier_model: Verifier,
vlm_model: Optional[Verifier] = None,
sinks: Optional[list[Sink]] = None,
inference_interval: float = 0.5,
active_interval: float = 0.0,
cooldown_frames: int = 60,
config: Optional[PipelineConfig] = None,
):
self.person_model = person_model
self.cigarette_model = cigarette_model
self.verifier_model = verifier_model
self.vlm_model = vlm_model
self.sinks = sinks or []
self.cfg = config or DEFAULT_CONFIG
# Adaptive Sampling Settings
self.idle_interval = inference_interval
self.active_interval = active_interval
self.cooldown_frames = cooldown_frames
# State
self.current_interval = self.idle_interval
# Adaptive sampling state
self.current_interval = self.cfg.idle_interval
self.active_cooldown = 0
self.last_inference_time = 0.0
# Async Inference State
# Async inference state
self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
self.current_future: Optional[Future] = None
self.latest_detections: List[
Detection
] = [] # Cache for display to prevent flickering
self.latest_ticket: Optional[Ticket] = None
self.current_future: Optional[Future[InferenceResult]] = None
self.latest_result: Optional[InferenceResult] = None
# Buffer for temporal context (optional, for VLM or sequence logic)
# Frame buffer for temporal context
self.frame_buffer = deque(maxlen=30)
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):
"""Main loop processing the video stream."""
"""Main loop: read frames, schedule inference, dispatch to sinks."""
logger.info("Starting stream processing...")
count = 0
frame_count = 0
for frame in source.stream():
count += 1
if count % 100 == 0:
logger.info(f"Processed {count} frames...")
frame_count += 1
if frame_count % 100 == 0:
logger.info("Processed %d frames...", frame_count)
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
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.")
logger.info("Stream finished. Total frames: %d", frame_count)
self.executor.shutdown(wait=False)
def _run_inference(
self, frame: Frame
) -> Tuple[bool, List[Detection], Optional[Ticket]]:
def _poll_inference_result(self, frame: Frame):
"""Check if async inference completed, update state."""
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.
Returns (ActivityDetected, ListOfDetections, TicketOrNone).
Does NOT modify frame in place (thread safety).
Full detection chain. Thread-safe (doesn't mutate frame).
Returns InferenceResult with detections and optional ticket.
"""
t0 = time.time()
detections = []
# --- Step 1: Broad Phase (Person Detection) ---
people = self.person_model.predict(frame.data)
people = nms_by_class(people, iou_threshold=0.7)
detections: List[Detection] = []
# Stage 1: Person Detection
people = self._detect_people(frame)
detections.extend(people)
if not people:
# logger.debug("No people found.")
return False, detections, None
return InferenceResult(has_activity=False, detections=detections)
# --- Step 2: Narrow Phase (Cigarette Candidate) ---
candidates = self._find_candidates(frame, people, detections)
# Stage 2: Cigarette Detection (with early exit for high confidence)
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:
return True, detections, None
return InferenceResult(has_activity=True, detections=detections)
# --- Step 3: Verification (Classifier) ---
best_hit = self._verify_candidates(frame, candidates, detections)
# Stage 3: Verification
verified_hit = self._verify_candidates(frame, candidates, detections)
ticket = None
# --- Step 4: Action ---
if best_hit:
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",
)
if verified_hit:
ticket = self._create_ticket(frame, verified_hit, stage="Verifier")
# Add the verified box to detections for visualization
verified_det = Detection(
box=BoundingBox(
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)
detections = nms_by_class(detections, self.cfg.final_nms_iou)
self._log_inference_time(t0)
return InferenceResult(has_activity=True, detections=detections, ticket=ticket)
# Clean up overlapping boxes from model output
detections = nms_by_class(detections, iou_threshold=0.5)
def _detect_people(self, frame: Frame) -> List[Detection]:
"""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
if dt > 0.1:
logger.debug("Inference took %.3fs (Slow)", dt)
def _detect_cigarettes(
self,
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
def _find_candidates(
self, frame: Frame, people: List[Detection], detections_out: List[Detection]
) -> 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)
Returns:
(candidates, early_hit)
- candidates: list of (person, cigs, offset) for verification
- early_hit: (person, cig, score, offset) if high-confidence detection, else None
"""
candidates = []
early_hit = None
h_img, w_img = frame.data.shape[:2]
margin = 0.2 # 20% expansion
cfg = self.cfg
for person in people:
# 1. Expand Box
pb = person.box
bw = pb.x2 - pb.x1
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
# Expand person box
crop_box = self._expand_box(person.box, cfg.person_margin, w_img, h_img)
ex1, ey1, ex2, ey2 = crop_box
person_crop = frame.data[ey1:ey2, ex1:ex2]
if (
person_crop.size == 0
or person_crop.shape[0] < 10
or person_crop.shape[1] < 10
):
if not self._is_valid_crop(person_crop):
continue
# 3. Detect
cig_candidates = self.cigarette_model.predict(person_crop)
# Detect cigarettes in 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)
cig_candidates = [
c
for c in cig_candidates
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)
for cig in cigs:
# Add to visualization
global_det = self._to_global_coords(cig, ex1, ey1, "candidate_cig")
detections_out.append(global_det)
if cig_candidates:
candidates.append((person, cig_candidates, (ex1, ey1)))
# 5. Global Coordinate Mapping (For Visualization)
for cig in cig_candidates:
global_cig = Detection(
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,
# Check for early exit (high confidence)
if cig.score >= cfg.cig_high_confidence and early_hit is None:
early_hit = (person, cig, cig.score, (ex1, ey1))
logger.info(
"High-confidence cig detection (%.2f) - skipping verifier",
cig.score,
)
detections_out.append(global_cig)
return candidates
if cigs:
candidates.append((person, cigs, (ex1, ey1)))
return candidates, early_hit
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]:
"""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 cig in cigs:
# 1. Derive Crop Coordinates
cig_w = cig.box.x2 - cig.box.x1
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:
# Already caught by early exit?
if cig.score >= cfg.cig_high_confidence:
continue
# 3. Verify
score = self.verifier_model.verify(obj_crop)
# Extract padded crop for classifier
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
confirmed_hits.append((person, cig, score))
score = self.verifier_model.verify(crop)
verified_det = Detection(
box=BoundingBox(
x1=ex1 + cig.box.x1,
y1=ey1 + cig.box.y1,
x2=ex1 + cig.box.x2,
y2=ey1 + cig.box.y2,
conf=score,
),
score=score,
label="verified_smoking",
metadata={"verifier_score": score},
if score >= cfg.verify_threshold:
confirmed.append((person, cig, score, (ex1, ey1)))
# Add verified detection
verified_det = self._to_global_coords(
cig, ex1, ey1, "verified_smoking"
)
verified_det.score = score
verified_det.metadata["verifier_score"] = score
detections_out.append(verified_det)
if confirmed_hits:
return max(confirmed_hits, key=lambda x: x[2])
return None
return max(confirmed, key=lambda x: x[2]) if confirmed else 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)