This commit is contained in:
lda committed 2026-01-22 21:11:10 +07:00
1 parent 73b1705fa3
commit 792eda9ab3
10 files changed
+1171 -13

No files matched your search

+1 -1
View File
@@ -33,7 +33,7 @@ class ThreadedStreamSource(StreamSource):
if cap.isOpened():
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
cap.release()
except cv2.error: # pylint: disable=catching-non-exception
except cv2.error: # pylint: disable=catching-non-exception
logger.warning("Could not probe FPS from source.")
self.stopped = False
+190
View File
@@ -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,
)
+26
View File
@@ -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",
]
+372
View File
@@ -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"}
+257
View File
@@ -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
+109
View File
@@ -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.)."""
...
+86
View File
@@ -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)