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

+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)