ahh
This commit is contained in:
1 parent
73b1705fa3
commit
792eda9ab3
9 files changed
+1162
-4
No files matched your search
@@ -94,6 +94,11 @@ uv run -m tools.webcam_check
|
||||
| `--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. |
|
||||
| `--staged` | `False` | Use stage-based pipeline (modular, supports VLM). |
|
||||
| `--vlm` | `False` | Enable VLM verification (GPT-4V, Claude Vision). |
|
||||
| `--vlm-key` | env var | API key for VLM (or set `OPENAI_API_KEY` env var). |
|
||||
| `--vlm-provider` | `openai` | VLM provider: `openai` or `anthropic`. |
|
||||
| `--vlm-model` | `gpt-4-vision-preview` | VLM model name. |
|
||||
|
||||
## 📊 Outputs
|
||||
|
||||
@@ -147,7 +152,8 @@ gp/
|
||||
│
|
||||
├── pipeline/
|
||||
│ ├── config.py # PipelineConfig dataclass (all thresholds)
|
||||
│ ├── core.py # SmokingDetectionPipeline (async inference)
|
||||
│ ├── core.py # SmokingDetectionPipeline (classic, fixed stages)
|
||||
│ ├── staged_core.py # StageBasedPipeline (modular, VLM support)
|
||||
│ ├── protocols.py # Interfaces (Model, Verifier, Sink, StreamSource)
|
||||
│ ├── types.py # Data classes (Frame, Detection, Ticket, BoundingBox)
|
||||
│ │
|
||||
@@ -166,6 +172,12 @@ gp/
|
||||
│ │ ├── screenshots.py # ScreenshotSink (evidence images)
|
||||
│ │ └── video.py # VideoFileStore (async video writer)
|
||||
│ │
|
||||
│ ├── stages/ # Pluggable stage system
|
||||
│ │ ├── protocol.py # Stage protocol, StageContext, StageResult
|
||||
│ │ ├── runner.py # StagePipeline (executes stage chain)
|
||||
│ │ ├── local.py # LocalPersonDetector, LocalCigaretteDetector, LocalVerifier
|
||||
│ │ └── api.py # APIPersonDetector, VLMVerifier (cloud stages)
|
||||
│ │
|
||||
│ └── ops/ # Algorithms & utilities
|
||||
│ ├── nms.py # Non-Maximum Suppression
|
||||
│ ├── geometry.py # Box expansion, zone filtering
|
||||
@@ -219,3 +231,54 @@ 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
|
||||
|
||||
## 🔌 Pluggable Stage System
|
||||
|
||||
The `--staged` flag enables the modular stage pipeline, allowing you to mix local models with cloud APIs:
|
||||
|
||||
### Stage Types
|
||||
|
||||
| Stage | Description | Use Case |
|
||||
|:------|:------------|:---------|
|
||||
| `LocalPersonDetector` | YOLO person detection | Default, GPU/CPU |
|
||||
| `LocalCigaretteDetector` | YOLO cigarette detection | Default, GPU/CPU |
|
||||
| `LocalVerifier` | YOLO classifier | Default, GPU/CPU |
|
||||
| `APIPersonDetector` | Cloud person detection | No local GPU, scaling |
|
||||
| `VLMVerifier` | GPT-4V/Claude Vision | High-stakes verification |
|
||||
|
||||
### VLM Verification Example
|
||||
|
||||
```bash
|
||||
# Use GPT-4V for final verification
|
||||
export OPENAI_API_KEY="sk-..."
|
||||
uv run main.py --source 0 --display --vlm
|
||||
|
||||
# Use Claude Vision instead
|
||||
export ANTHROPIC_API_KEY="sk-ant-..."
|
||||
uv run main.py --source video.mp4 --vlm --vlm-provider anthropic --vlm-model claude-sonnet-4-20250514
|
||||
```
|
||||
|
||||
### Custom Stage Configuration
|
||||
|
||||
```python
|
||||
from pipeline.stages import (
|
||||
LocalPersonDetector,
|
||||
LocalCigaretteDetector,
|
||||
VLMVerifier,
|
||||
StagePipeline,
|
||||
)
|
||||
|
||||
# Build custom stage chain
|
||||
stages = [
|
||||
LocalPersonDetector(person_model, config),
|
||||
LocalCigaretteDetector(cig_model, config),
|
||||
# Skip local verifier, go straight to VLM
|
||||
VLMVerifier(api_key="sk-...", provider="openai"),
|
||||
]
|
||||
|
||||
pipeline = StagePipeline(stages, config)
|
||||
ctx = pipeline.run(frame)
|
||||
|
||||
if ctx.ticket:
|
||||
print(f"Smoking detected! Confidence: {ctx.ticket.confidence}")
|
||||
```
|
||||
@@ -1,6 +1,7 @@
|
||||
import argparse
|
||||
import logging
|
||||
import logging.handlers
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
@@ -16,6 +17,7 @@ from pipeline.sinks import (
|
||||
WindowDisplaySink,
|
||||
)
|
||||
from pipeline.sources import FileStreamSource, ThreadedRTSPSource
|
||||
from pipeline.staged_core import StageBasedPipeline
|
||||
|
||||
|
||||
def setup_logging(debug: bool = False):
|
||||
@@ -99,6 +101,35 @@ def main():
|
||||
action="store_true",
|
||||
help="Don't save output video (useful with --display)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--staged",
|
||||
action="store_true",
|
||||
help="Use stage-based pipeline (allows VLM verification)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vlm",
|
||||
action="store_true",
|
||||
help="Enable VLM verification (requires --vlm-key or OPENAI_API_KEY env var)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vlm-key",
|
||||
type=str,
|
||||
default=None,
|
||||
help="API key for VLM (default: uses OPENAI_API_KEY env var)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vlm-provider",
|
||||
type=str,
|
||||
default="openai",
|
||||
choices=["openai", "anthropic"],
|
||||
help="VLM provider (default: openai)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vlm-model",
|
||||
type=str,
|
||||
default="gpt-4-vision-preview",
|
||||
help="VLM model name (default: gpt-4-vision-preview)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -173,6 +204,33 @@ def main():
|
||||
cooldown_frames=int(source.fps * 2), # Stay active 2s after detection
|
||||
)
|
||||
|
||||
# Choose pipeline implementation
|
||||
if args.staged or args.vlm:
|
||||
# Stage-based pipeline (modular, supports VLM)
|
||||
vlm_key = (
|
||||
args.vlm_key or os.environ.get("OPENAI_API_KEY")
|
||||
if args.vlm_provider == "openai"
|
||||
else os.environ.get("ANTHROPIC_API_KEY")
|
||||
)
|
||||
|
||||
if args.vlm and not vlm_key:
|
||||
logger.warning(
|
||||
"VLM enabled but no API key found. Set --vlm-key or OPENAI_API_KEY/ANTHROPIC_API_KEY env var."
|
||||
)
|
||||
|
||||
pipeline = StageBasedPipeline(
|
||||
person_model=person_model,
|
||||
cigarette_model=cigarette_model,
|
||||
verifier_model=verifier,
|
||||
vlm_api_key=vlm_key if args.vlm else None,
|
||||
vlm_provider=args.vlm_provider,
|
||||
vlm_model=args.vlm_model,
|
||||
use_vlm=args.vlm,
|
||||
sinks=sinks,
|
||||
config=config,
|
||||
)
|
||||
else:
|
||||
# Classic pipeline (fixed stages, battle-tested)
|
||||
pipeline = SmokingDetectionPipeline(
|
||||
person_model=person_model,
|
||||
cigarette_model=cigarette_model,
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Stage-based pipeline orchestration."""
|
||||
|
||||
import concurrent.futures
|
||||
import logging
|
||||
from collections import deque
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
|
||||
from pipeline.protocols import Model, Sink, StreamSource, Verifier
|
||||
from pipeline.stages import (
|
||||
LocalCigaretteDetector,
|
||||
LocalPersonDetector,
|
||||
LocalVerifier,
|
||||
StagePipeline,
|
||||
VLMVerifier,
|
||||
)
|
||||
from pipeline.types import Detection, Frame, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class InferenceResult:
|
||||
"""Result of a stage pipeline run."""
|
||||
|
||||
has_activity: bool
|
||||
detections: List[Detection]
|
||||
ticket: Optional[Ticket] = None
|
||||
|
||||
|
||||
class StageBasedPipeline:
|
||||
"""
|
||||
Modular smoking detection pipeline using pluggable stages.
|
||||
|
||||
This is the "next-gen" version of SmokingDetectionPipeline. Instead of
|
||||
hardcoded model calls, it uses configurable Stage objects.
|
||||
|
||||
Example configurations:
|
||||
|
||||
1. Full local (default):
|
||||
LocalPersonDetector -> LocalCigaretteDetector -> LocalVerifier
|
||||
|
||||
2. VLM verification:
|
||||
LocalPersonDetector -> LocalCigaretteDetector -> VLMVerifier
|
||||
|
||||
3. Hybrid (local + cloud fallback):
|
||||
LocalPersonDetector -> LocalCigaretteDetector -> LocalVerifier -> VLMVerifier
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
person_model: Model,
|
||||
cigarette_model: Model,
|
||||
verifier_model: Verifier,
|
||||
vlm_api_key: Optional[str] = None,
|
||||
vlm_provider: str = "openai",
|
||||
vlm_model: str = "gpt-4-vision-preview",
|
||||
use_vlm: bool = False,
|
||||
sinks: Optional[list[Sink]] = None,
|
||||
config: Optional[PipelineConfig] = None,
|
||||
):
|
||||
self.sinks = sinks or []
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
# Build stage pipeline
|
||||
stages = [
|
||||
LocalPersonDetector(person_model, config),
|
||||
LocalCigaretteDetector(cigarette_model, config),
|
||||
LocalVerifier(verifier_model, config),
|
||||
]
|
||||
|
||||
if use_vlm and vlm_api_key:
|
||||
stages.append(
|
||||
VLMVerifier(
|
||||
api_key=vlm_api_key,
|
||||
provider=vlm_provider,
|
||||
model=vlm_model,
|
||||
config=config,
|
||||
)
|
||||
)
|
||||
|
||||
self.stage_pipeline = StagePipeline(stages, config)
|
||||
|
||||
# Adaptive sampling state
|
||||
self.current_interval = self.cfg.idle_interval
|
||||
self.active_cooldown = 0
|
||||
self.last_inference_time = 0.0
|
||||
|
||||
# Async inference state
|
||||
self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
|
||||
self.current_future: Optional[Future[InferenceResult]] = None
|
||||
self.latest_result: Optional[InferenceResult] = None
|
||||
|
||||
# Frame buffer for temporal context
|
||||
self.frame_buffer = deque(maxlen=30)
|
||||
|
||||
logger.info(
|
||||
"StageBasedPipeline initialized with %d sinks, %d stages. Config: %s",
|
||||
len(self.sinks),
|
||||
len(stages),
|
||||
self.cfg,
|
||||
)
|
||||
|
||||
def process_stream(self, source: StreamSource):
|
||||
"""Main loop: read frames, schedule inference, dispatch to sinks."""
|
||||
logger.info("Starting stream processing (stage-based)...")
|
||||
frame_count = 0
|
||||
|
||||
for frame in source.stream():
|
||||
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)
|
||||
|
||||
logger.info("Stream finished. Total frames: %d", frame_count)
|
||||
self.executor.shutdown(wait=False)
|
||||
self.stage_pipeline.close()
|
||||
|
||||
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 = list(self.latest_result.detections)
|
||||
|
||||
# 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)
|
||||
|
||||
def _run_inference(self, frame: Frame) -> InferenceResult:
|
||||
"""Run stage pipeline on frame."""
|
||||
ctx = self.stage_pipeline.run(frame)
|
||||
|
||||
has_activity = len(ctx.people) > 0
|
||||
|
||||
return InferenceResult(
|
||||
has_activity=has_activity,
|
||||
detections=ctx.all_detections,
|
||||
ticket=ctx.ticket,
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Pluggable pipeline stages."""
|
||||
|
||||
from pipeline.stages.api import APIPersonDetector, VLMVerifier
|
||||
from pipeline.stages.local import (
|
||||
LocalCigaretteDetector,
|
||||
LocalPersonDetector,
|
||||
LocalVerifier,
|
||||
)
|
||||
from pipeline.stages.protocol import Stage, StageContext, StageResult
|
||||
from pipeline.stages.runner import StagePipeline
|
||||
|
||||
__all__ = [
|
||||
# Protocol
|
||||
"Stage",
|
||||
"StageContext",
|
||||
"StageResult",
|
||||
# Runner
|
||||
"StagePipeline",
|
||||
# Local stages
|
||||
"LocalPersonDetector",
|
||||
"LocalCigaretteDetector",
|
||||
"LocalVerifier",
|
||||
# API stages
|
||||
"APIPersonDetector",
|
||||
"VLMVerifier",
|
||||
]
|
||||
@@ -0,0 +1,372 @@
|
||||
"""API-based stages for cloud/VLM verification."""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import time
|
||||
from abc import ABC
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
|
||||
from pipeline.stages.protocol import StageContext, StageResult
|
||||
from pipeline.types import BoundingBox, Detection, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class APIStageBase(ABC):
|
||||
"""Base class for API-based stages with common HTTP/retry logic."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
endpoint: Optional[str] = None,
|
||||
timeout: float = 30.0,
|
||||
max_retries: int = 3,
|
||||
config: Optional[PipelineConfig] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.endpoint = endpoint
|
||||
self.timeout = timeout
|
||||
self.max_retries = max_retries
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
def _encode_image(self, image: np.ndarray) -> str:
|
||||
"""Encode image as base64 for API calls."""
|
||||
import cv2
|
||||
|
||||
_, buffer = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||
return base64.b64encode(buffer).decode("utf-8")
|
||||
|
||||
def _make_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Make HTTP request with retries. Override for specific APIs."""
|
||||
raise NotImplementedError("Subclass must implement _make_request")
|
||||
|
||||
def close(self):
|
||||
"""Release any API connections/sessions."""
|
||||
pass
|
||||
|
||||
|
||||
class APIPersonDetector(APIStageBase):
|
||||
"""
|
||||
Stage 1: Person detection via cloud API.
|
||||
|
||||
Useful when:
|
||||
- No GPU available locally
|
||||
- Need to scale inference
|
||||
- Using specialized person detection service
|
||||
|
||||
Compatible APIs:
|
||||
- AWS Rekognition
|
||||
- Google Cloud Vision
|
||||
- Azure Computer Vision
|
||||
- Roboflow Hosted Inference
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
endpoint: Optional[str] = None,
|
||||
provider: str = "generic",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(api_key=api_key, endpoint=endpoint, **kwargs)
|
||||
self.provider = provider
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"api_person_detector_{self.provider}"
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
if not self.api_key:
|
||||
return StageResult.fail("API key not configured")
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
try:
|
||||
image_b64 = self._encode_image(ctx.frame_data)
|
||||
response = self._make_request(
|
||||
{"image": image_b64, "task": "person_detection"}
|
||||
)
|
||||
|
||||
people = self._parse_response(response, ctx.frame_data.shape)
|
||||
ctx.people = people
|
||||
ctx.all_detections.extend(people)
|
||||
|
||||
ctx.stage_times[self.name] = time.time() - t0
|
||||
|
||||
if not people:
|
||||
ctx.signal_stop("no_people")
|
||||
return StageResult.skip("No people detected (API)")
|
||||
|
||||
return StageResult.ok(people, api_latency=time.time() - t0)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("API person detection failed: %s", e)
|
||||
return StageResult.fail(str(e))
|
||||
|
||||
def _parse_response(self, response: Dict, img_shape) -> List[Detection]:
|
||||
"""Parse API response into Detection objects. Override per provider."""
|
||||
# Generic format: {"detections": [{"box": [x1,y1,x2,y2], "confidence": 0.9}]}
|
||||
detections = []
|
||||
h, w = img_shape[:2]
|
||||
|
||||
for det in response.get("detections", []):
|
||||
box = det.get("box", [0, 0, 0, 0])
|
||||
conf = det.get("confidence", 0.0)
|
||||
|
||||
detections.append(
|
||||
Detection(
|
||||
box=BoundingBox(
|
||||
x1=int(box[0] * w) if box[0] <= 1 else int(box[0]),
|
||||
y1=int(box[1] * h) if box[1] <= 1 else int(box[1]),
|
||||
x2=int(box[2] * w) if box[2] <= 1 else int(box[2]),
|
||||
y2=int(box[3] * h) if box[3] <= 1 else int(box[3]),
|
||||
conf=conf,
|
||||
),
|
||||
score=conf,
|
||||
label="person",
|
||||
metadata={"source": "api", "provider": self.provider},
|
||||
)
|
||||
)
|
||||
|
||||
return detections
|
||||
|
||||
def _make_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Override with actual API call."""
|
||||
# Placeholder - implement for specific provider
|
||||
import requests
|
||||
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
resp = requests.post(
|
||||
self.endpoint,
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
class VLMVerifier(APIStageBase):
|
||||
"""
|
||||
Vision-Language Model verification stage.
|
||||
|
||||
Uses GPT-4V, Claude Vision, or similar to verify smoking.
|
||||
More accurate than local classifiers, handles edge cases.
|
||||
|
||||
Best for:
|
||||
- Final verification of high-stakes alerts
|
||||
- Handling unusual scenarios
|
||||
- When false positives are costly
|
||||
"""
|
||||
|
||||
SYSTEM_PROMPT = """You are a smoking detection verification system.
|
||||
Analyze the image and determine if someone is actively smoking (cigarette, vape, etc.).
|
||||
|
||||
Respond with JSON only:
|
||||
{
|
||||
"is_smoking": true/false,
|
||||
"confidence": 0.0-1.0,
|
||||
"reasoning": "brief explanation"
|
||||
}
|
||||
|
||||
Be conservative - only confirm if clearly visible."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
model: str = "gpt-4-vision-preview",
|
||||
provider: str = "openai",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(api_key=api_key, **kwargs)
|
||||
self.model = model
|
||||
self.provider = provider
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"vlm_verifier_{self.provider}"
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
if not self.api_key:
|
||||
return StageResult.fail("VLM API key not configured")
|
||||
|
||||
if not ctx.cigarettes:
|
||||
return StageResult.skip("No cigarettes to verify")
|
||||
|
||||
if ctx.should_stop and ctx.stop_reason == "high_confidence_early_exit":
|
||||
return StageResult.skip("Already confirmed by detector")
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
try:
|
||||
# Use full frame or crop around detections
|
||||
image = self._prepare_verification_image(ctx)
|
||||
result = self._call_vlm(image)
|
||||
|
||||
ctx.stage_times[self.name] = time.time() - t0
|
||||
|
||||
if (
|
||||
result.get("is_smoking")
|
||||
and result.get("confidence", 0) >= self.cfg.verify_threshold
|
||||
):
|
||||
score = result["confidence"]
|
||||
|
||||
# Mark best cigarette as verified
|
||||
if ctx.cigarettes:
|
||||
best_cig = max(ctx.cigarettes, key=lambda d: d.score)
|
||||
verified_det = Detection(
|
||||
box=best_cig.box,
|
||||
score=score,
|
||||
label="vlm_verified_smoking",
|
||||
metadata={
|
||||
"vlm_reasoning": result.get("reasoning", ""),
|
||||
"vlm_model": self.model,
|
||||
},
|
||||
)
|
||||
ctx.verified.append(verified_det)
|
||||
ctx.all_detections.append(verified_det)
|
||||
|
||||
ctx.ticket = Ticket(
|
||||
frame_id=ctx.frame_id,
|
||||
confidence=score,
|
||||
evidence=[ctx.frame_data],
|
||||
stage_confirmed=f"VLM:{self.model}",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"VLM CONFIRMED SMOKING! Score: %.2f, Reason: %s",
|
||||
score,
|
||||
result.get("reasoning", ""),
|
||||
)
|
||||
|
||||
return StageResult.ok(ctx.verified, vlm_confidence=score)
|
||||
|
||||
return StageResult.skip("VLM did not confirm smoking")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("VLM verification failed: %s", e)
|
||||
return StageResult.fail(str(e))
|
||||
|
||||
def _prepare_verification_image(self, ctx: StageContext) -> np.ndarray:
|
||||
"""Prepare image for VLM - either full frame or focused crop."""
|
||||
# For now, use full frame. Could crop to detection area.
|
||||
return ctx.frame_data
|
||||
|
||||
def _call_vlm(self, image: np.ndarray) -> Dict[str, Any]:
|
||||
"""Call VLM API. Override for specific providers."""
|
||||
image_b64 = self._encode_image(image)
|
||||
|
||||
if self.provider == "openai":
|
||||
return self._call_openai(image_b64)
|
||||
elif self.provider == "anthropic":
|
||||
return self._call_anthropic(image_b64)
|
||||
else:
|
||||
raise ValueError(f"Unknown VLM provider: {self.provider}")
|
||||
|
||||
def _call_openai(self, image_b64: str) -> Dict[str, Any]:
|
||||
"""Call OpenAI GPT-4V."""
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{"role": "system", "content": self.SYSTEM_PROMPT},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Is someone smoking in this image?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
"max_tokens": 200,
|
||||
}
|
||||
|
||||
resp = requests.post(
|
||||
"https://api.openai.com/v1/chat/completions",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
|
||||
# Parse JSON from response
|
||||
try:
|
||||
# Handle markdown code blocks
|
||||
if "```json" in content:
|
||||
content = content.split("```json")[1].split("```")[0]
|
||||
elif "```" in content:
|
||||
content = content.split("```")[1].split("```")[0]
|
||||
return json.loads(content.strip())
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Failed to parse VLM response: %s", content)
|
||||
return {"is_smoking": False, "confidence": 0, "reasoning": "Parse error"}
|
||||
|
||||
def _call_anthropic(self, image_b64: str) -> Dict[str, Any]:
|
||||
"""Call Anthropic Claude Vision."""
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
headers = {
|
||||
"x-api-key": self.api_key,
|
||||
"Content-Type": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"max_tokens": 200,
|
||||
"system": self.SYSTEM_PROMPT,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/jpeg",
|
||||
"data": image_b64,
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "Is someone smoking in this image?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
resp = requests.post(
|
||||
"https://api.anthropic.com/v1/messages",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
content = resp.json()["content"][0]["text"]
|
||||
|
||||
try:
|
||||
if "```json" in content:
|
||||
content = content.split("```json")[1].split("```")[0]
|
||||
elif "```" in content:
|
||||
content = content.split("```")[1].split("```")[0]
|
||||
return json.loads(content.strip())
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Failed to parse VLM response: %s", content)
|
||||
return {"is_smoking": False, "confidence": 0, "reasoning": "Parse error"}
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Local model-based stages using YOLO."""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
|
||||
from pipeline.ops import nms_by_class
|
||||
from pipeline.protocols import Model, Verifier
|
||||
from pipeline.stages.protocol import StageContext, StageResult
|
||||
from pipeline.types import BoundingBox, Detection, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LocalPersonDetector:
|
||||
"""Stage 1: Detect people using local YOLO model."""
|
||||
|
||||
def __init__(self, model: Model, config: Optional[PipelineConfig] = None):
|
||||
self.model = model
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "local_person_detector"
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
t0 = time.time()
|
||||
|
||||
people = self.model.predict(ctx.frame_data)
|
||||
people = nms_by_class(people, self.cfg.person_nms_iou)
|
||||
|
||||
# Write to context
|
||||
ctx.people = people
|
||||
ctx.all_detections.extend(people)
|
||||
|
||||
ctx.stage_times[self.name] = time.time() - t0
|
||||
|
||||
if not people:
|
||||
ctx.signal_stop("no_people")
|
||||
return StageResult.skip("No people detected")
|
||||
|
||||
return StageResult.ok(people)
|
||||
|
||||
def close(self):
|
||||
if hasattr(self.model, "close"):
|
||||
self.model.close() # pyright: ignore
|
||||
|
||||
|
||||
class LocalCigaretteDetector:
|
||||
"""Stage 2: Detect cigarettes in person crops using local YOLO model."""
|
||||
|
||||
def __init__(self, model: Model, config: Optional[PipelineConfig] = None):
|
||||
self.model = model
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "local_cigarette_detector"
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
if not ctx.people:
|
||||
return StageResult.skip("No people to scan")
|
||||
|
||||
t0 = time.time()
|
||||
cfg = self.cfg
|
||||
h_img, w_img = ctx.frame_data.shape[:2]
|
||||
|
||||
all_cigs: List[Detection] = []
|
||||
early_hit = None
|
||||
|
||||
for person in ctx.people:
|
||||
# Expand person box for crop
|
||||
crop_box = self._expand_box(person.box, cfg.person_margin, w_img, h_img)
|
||||
ex1, ey1, ex2, ey2 = crop_box
|
||||
person_crop = ctx.frame_data[ey1:ey2, ex1:ex2]
|
||||
|
||||
if not self._is_valid_crop(person_crop):
|
||||
continue
|
||||
|
||||
# Detect cigarettes in crop
|
||||
cigs = self.model.predict(person_crop)
|
||||
cigs = self._filter_small_boxes(cigs, cfg.min_cig_size)
|
||||
cigs = nms_by_class(cigs, cfg.cig_nms_iou)
|
||||
|
||||
for cig in cigs:
|
||||
# Convert to global coords
|
||||
global_cig = self._to_global_coords(cig, ex1, ey1, "candidate_cig")
|
||||
global_cig.metadata["person_box"] = (
|
||||
person.box.x1,
|
||||
person.box.y1,
|
||||
person.box.x2,
|
||||
person.box.y2,
|
||||
)
|
||||
global_cig.metadata["crop_offset"] = (ex1, ey1)
|
||||
all_cigs.append(global_cig)
|
||||
ctx.all_detections.append(global_cig)
|
||||
|
||||
# Check for early exit (high confidence)
|
||||
if cig.score >= cfg.cig_high_confidence and early_hit is None:
|
||||
early_hit = {
|
||||
"person": person,
|
||||
"cig": cig,
|
||||
"global_cig": global_cig,
|
||||
"score": cig.score,
|
||||
"offset": (ex1, ey1),
|
||||
}
|
||||
|
||||
ctx.cigarettes = all_cigs
|
||||
ctx.stage_times[self.name] = time.time() - t0
|
||||
|
||||
# Handle early exit for high-confidence detection
|
||||
if early_hit:
|
||||
logger.info(
|
||||
"High-confidence cig detection (%.2f) - skipping verifier",
|
||||
early_hit["score"],
|
||||
)
|
||||
ctx.verified.append(early_hit["global_cig"])
|
||||
ctx.ticket = Ticket(
|
||||
frame_id=ctx.frame_id,
|
||||
confidence=early_hit["score"],
|
||||
evidence=[ctx.frame_data],
|
||||
stage_confirmed="Detector",
|
||||
)
|
||||
ctx.signal_stop("high_confidence_early_exit")
|
||||
return StageResult.ok(all_cigs, early_exit=True, score=early_hit["score"])
|
||||
|
||||
if not all_cigs:
|
||||
return StageResult.skip("No cigarettes detected")
|
||||
|
||||
return StageResult.ok(all_cigs)
|
||||
|
||||
def _expand_box(self, box: BoundingBox, margin: float, w_max: int, h_max: int):
|
||||
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:
|
||||
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):
|
||||
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
|
||||
):
|
||||
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 close(self):
|
||||
if hasattr(self.model, "close"):
|
||||
self.model.close() # pyright: ignore
|
||||
|
||||
|
||||
class LocalVerifier:
|
||||
"""Stage 3: Verify cigarette detections using local classifier."""
|
||||
|
||||
def __init__(self, model: Verifier, config: Optional[PipelineConfig] = None):
|
||||
self.model = model
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "local_verifier"
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
if not ctx.cigarettes:
|
||||
return StageResult.skip("No cigarettes to verify")
|
||||
|
||||
if ctx.should_stop:
|
||||
return StageResult.skip(f"Pipeline stopped: {ctx.stop_reason}")
|
||||
|
||||
t0 = time.time()
|
||||
cfg = self.cfg
|
||||
h_img, w_img = ctx.frame_data.shape[:2]
|
||||
|
||||
confirmed: list[tuple[Detection, float]] = []
|
||||
|
||||
for cig_det in ctx.cigarettes:
|
||||
offset = cig_det.metadata.get("crop_offset", (0, 0))
|
||||
ex1, ey1 = offset
|
||||
|
||||
# Extract padded crop for classifier
|
||||
crop = self._extract_cig_crop(ctx.frame_data, cig_det, w_img, h_img)
|
||||
if crop.size == 0:
|
||||
continue
|
||||
|
||||
score = self.model.verify(crop)
|
||||
|
||||
if score >= cfg.verify_threshold:
|
||||
verified_det = Detection(
|
||||
box=cig_det.box,
|
||||
score=score,
|
||||
label="verified_smoking",
|
||||
metadata={**cig_det.metadata, "verifier_score": score},
|
||||
)
|
||||
confirmed.append((verified_det, score))
|
||||
ctx.verified.append(verified_det)
|
||||
ctx.all_detections.append(verified_det)
|
||||
|
||||
ctx.stage_times[self.name] = time.time() - t0
|
||||
|
||||
if confirmed:
|
||||
# Create ticket from best detection
|
||||
best_det, best_score = max(confirmed, key=lambda x: x[1])
|
||||
ctx.ticket = Ticket(
|
||||
frame_id=ctx.frame_id,
|
||||
confidence=best_score,
|
||||
evidence=[ctx.frame_data],
|
||||
stage_confirmed="Verifier",
|
||||
)
|
||||
logger.info("CONFIRMED SMOKING! Score: %.2f (Stage: Verifier)", best_score)
|
||||
return StageResult.ok(
|
||||
[d for d, _ in confirmed], verified_count=len(confirmed)
|
||||
)
|
||||
|
||||
return StageResult.skip("No detections verified")
|
||||
|
||||
def _extract_cig_crop(self, image, cig_det: Detection, w_max: int, h_max: int):
|
||||
"""Extract padded crop around cigarette for verification."""
|
||||
box = cig_det.box
|
||||
pad = self.cfg.cig_crop_padding
|
||||
cig_w = box.x2 - box.x1
|
||||
cig_h = box.y2 - box.y1
|
||||
|
||||
# Global coords with padding
|
||||
x1 = max(0, int(box.x1 - cig_w * pad))
|
||||
y1 = max(0, int(box.y1 - cig_h * pad))
|
||||
x2 = min(w_max, int(box.x2 + cig_w * pad))
|
||||
y2 = min(h_max, int(box.y2 + cig_h * pad))
|
||||
|
||||
return image[y1:y2, x1:x2]
|
||||
|
||||
def close(self):
|
||||
if hasattr(self.model, "close"):
|
||||
self.model.close() # pyright: ignore
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Stage protocol and data structures for pluggable pipeline."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
|
||||
|
||||
import numpy as np
|
||||
|
||||
from pipeline.types import Detection, Ticket
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageContext:
|
||||
"""
|
||||
Data passed between pipeline stages.
|
||||
|
||||
Each stage reads from context, processes, and writes results back.
|
||||
This enables loose coupling - stages don't know about each other.
|
||||
"""
|
||||
|
||||
# Input
|
||||
frame_data: np.ndarray
|
||||
frame_id: int
|
||||
timestamp: float
|
||||
|
||||
# Stage outputs (accumulated)
|
||||
people: List[Detection] = field(default_factory=list)
|
||||
cigarettes: List[Detection] = field(default_factory=list)
|
||||
verified: List[Detection] = field(default_factory=list)
|
||||
|
||||
# Final outputs
|
||||
all_detections: List[Detection] = field(default_factory=list)
|
||||
ticket: Optional[Ticket] = None
|
||||
|
||||
# Control flow
|
||||
should_stop: bool = False # Early exit signal
|
||||
stop_reason: str = ""
|
||||
|
||||
# Metadata for debugging/logging
|
||||
stage_times: Dict[str, float] = field(default_factory=dict)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def signal_stop(self, reason: str):
|
||||
"""Signal early pipeline exit."""
|
||||
self.should_stop = True
|
||||
self.stop_reason = reason
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageResult:
|
||||
"""Result returned by a stage."""
|
||||
|
||||
success: bool
|
||||
detections: List[Detection] = field(default_factory=list)
|
||||
error: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def ok(
|
||||
cls, detections: Optional[List[Detection]] = None, **metadata
|
||||
) -> "StageResult":
|
||||
return cls(success=True, detections=detections or [], metadata=metadata)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, error: str) -> "StageResult":
|
||||
return cls(success=False, error=error)
|
||||
|
||||
@classmethod
|
||||
def skip(cls, reason: str) -> "StageResult":
|
||||
"""Stage skipped (not an error, just nothing to do)."""
|
||||
return cls(success=True, metadata={"skipped": True, "reason": reason})
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Stage(Protocol):
|
||||
"""
|
||||
Protocol for a pipeline stage.
|
||||
|
||||
Stages are pluggable processing units. Each stage:
|
||||
- Reads what it needs from StageContext
|
||||
- Processes (local model, API call, etc.)
|
||||
- Writes results back to StageContext
|
||||
- Optionally signals early exit
|
||||
|
||||
Examples:
|
||||
- LocalPersonDetector: Runs YOLO locally
|
||||
- APIPersonDetector: Calls cloud API
|
||||
- VLMVerifier: Uses GPT-4V for verification
|
||||
"""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Unique stage identifier."""
|
||||
...
|
||||
|
||||
def process(self, ctx: StageContext) -> StageResult:
|
||||
"""
|
||||
Process the context.
|
||||
|
||||
Args:
|
||||
ctx: Shared context with frame data and accumulated results
|
||||
|
||||
Returns:
|
||||
StageResult indicating success/failure and any detections
|
||||
"""
|
||||
...
|
||||
|
||||
def close(self):
|
||||
"""Release resources (models, connections, etc.)."""
|
||||
...
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Stage-based pipeline runner."""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from pipeline.config import DEFAULT_CONFIG, PipelineConfig
|
||||
from pipeline.ops import nms_by_class
|
||||
from pipeline.stages.protocol import Stage, StageContext
|
||||
from pipeline.types import Frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StagePipeline:
|
||||
"""
|
||||
Configurable stage-based pipeline.
|
||||
|
||||
Unlike SmokingDetectionPipeline (which has fixed stages), this allows
|
||||
mixing and matching stages:
|
||||
|
||||
Example configurations:
|
||||
- Full local: LocalPersonDetector -> LocalCigaretteDetector -> LocalVerifier
|
||||
- API verify: LocalPersonDetector -> LocalCigaretteDetector -> VLMVerifier
|
||||
- Full cloud: APIPersonDetector -> APICigaretteDetector -> VLMVerifier
|
||||
- Hybrid: LocalPersonDetector -> LocalCigaretteDetector -> LocalVerifier + VLMVerifier
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stages: List[Stage],
|
||||
config: Optional[PipelineConfig] = None,
|
||||
):
|
||||
self.stages = stages
|
||||
self.cfg = config or DEFAULT_CONFIG
|
||||
|
||||
stage_names = [s.name for s in stages]
|
||||
logger.info(
|
||||
"StagePipeline initialized with %d stages: %s", len(stages), stage_names
|
||||
)
|
||||
|
||||
def run(self, frame: Frame) -> StageContext:
|
||||
"""
|
||||
Run all stages on a frame.
|
||||
|
||||
Returns StageContext with accumulated results.
|
||||
"""
|
||||
ctx = StageContext(
|
||||
frame_data=frame.data,
|
||||
frame_id=frame.frame_id,
|
||||
timestamp=frame.timestamp,
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
for stage in self.stages:
|
||||
if ctx.should_stop:
|
||||
logger.debug(
|
||||
"Skipping stage '%s' - pipeline stopped: %s",
|
||||
stage.name,
|
||||
ctx.stop_reason,
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
result = stage.process(ctx)
|
||||
|
||||
if not result.success:
|
||||
logger.warning("Stage '%s' failed: %s", stage.name, result.error)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Stage '%s' raised exception: %s", stage.name, e)
|
||||
|
||||
# Apply final NMS across all detections
|
||||
ctx.all_detections = nms_by_class(ctx.all_detections, self.cfg.final_nms_iou)
|
||||
|
||||
ctx.metadata["total_time"] = time.time() - t0
|
||||
return ctx
|
||||
|
||||
def close(self):
|
||||
"""Release all stage resources."""
|
||||
for stage in self.stages:
|
||||
try:
|
||||
stage.close()
|
||||
except Exception as e:
|
||||
logger.error("Error closing stage '%s': %s", stage.name, e)
|
||||
LFS
BIN
Binary file not shown.
Reference in new issue
Block a user