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

I have a gpu! I want to use it!

Browse Source
This commit is contained in:
lda committed 2026-01-11 00:22:06 +07:00
1 parent b5ea0232f7
commit 013c082447
12 files changed
+252 -250
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pipeline/core.py
+64 -132
View File
@@ -1,20 +1,13 @@
import logging
from collections import deque
from typing import Optional
from typing import List, Optional, Tuple
from pipeline.protocols import AlertStore, Model, StreamSource, Verifier
from pipeline.types import Frame, Ticket
from pipeline.protocols import Model, Sink, StreamSource, Verifier
from pipeline.types import BoundingBox, Detection, Frame, Ticket
logger = logging.getLogger(__name__)
class ConsoleStore:
def store(self, ticket: Ticket):
logger.info(
f"[STORE] Persisting Alert! Conf: {ticket.confidence:.2f} @ Stage: {ticket.stage_confirmed}"
)
class SmokingDetectionPipeline:
def __init__(
self,
@@ -22,19 +15,20 @@ class SmokingDetectionPipeline:
cigarette_model: Model,
verifier_model: Verifier,
vlm_model: Optional[Verifier] = None,
store: Optional[AlertStore] = None,
sinks: Optional[list] = None,
sinks: Optional[list[Sink]] = None,
inference_interval: float = 0.5,
):
self.person_model = person_model
self.cigarette_model = cigarette_model
self.verifier_model = verifier_model
self.vlm_model = vlm_model
self.store = store or ConsoleStore()
self.sinks = sinks or []
self.inference_interval = inference_interval
self.last_inference_time = 0.0
# Buffer for temporal context (optional, for VLM or sequence logic)
self.frame_buffer = deque(maxlen=30)
logger.info("Pipeline initialized.")
logger.info(f"Pipeline initialized with {len(self.sinks)} sinks.")
def process_stream(self, source: StreamSource):
"""Main loop processing the video stream."""
@@ -47,7 +41,11 @@ class SmokingDetectionPipeline:
if count % 100 == 0:
logger.info(f"Processed {count} frames...")
self.frame_buffer.append(frame)
self._process_frame(frame)
# Skip inference if too soon
if (frame.timestamp - self.last_inference_time) >= self.inference_interval:
self._process_frame(frame)
self.last_inference_time = frame.timestamp
# Update sinks (Visualizers, etc)
for sink in self.sinks:
@@ -57,11 +55,7 @@ class SmokingDetectionPipeline:
def _process_frame(self, frame: Frame):
# --- Step 1: Broad Phase (Person Detection) ---
# "Is there a human here?"
# optimization: Resize or skip frames if inference is slow
people = self.person_model.predict(frame.data)
# Add to frame detections for visualization
frame.detections.extend(people)
if not people:
@@ -70,29 +64,49 @@ class SmokingDetectionPipeline:
logger.debug(f"Frame {frame.frame_id}: Found {len(people)} people.")
# --- Step 2: Narrow Phase (Cigarette Candidate) ---
# "Is the human holding something small?"
# We could optimize by only running on crops of the 'people'
candidates = []
for person in people:
# ROI: Region of Interest (e.g., upper body or hands)
# EXPANSTION: Expand the person box to include extended arms/surroundings
h_img, w_img = frame.data.shape[:2]
margin = 0.2 # 20% expansion
candidates = self._find_candidates(frame, people)
if not candidates:
return # EXIT
logger.debug(f"Frame {frame.frame_id}: Found {len(candidates)} candidates.")
# --- Step 3: Verification (Classifier) ---
best_hit = self._verify_candidates(frame, candidates)
# --- Step 4: Action ---
if best_hit:
person, cig, score = best_hit
ticket = Ticket(
frame_id=frame.frame_id,
confidence=score,
evidence=[frame.data],
stage_confirmed="Verifier", # Simplified for now
)
frame.ticket = ticket
def _find_candidates(
self, frame: Frame, people: List[Detection]
) -> List[Tuple[Detection, List[Detection], Tuple[int, int]]]:
"""Runs cigarette detection on expanded person crops."""
candidates = []
h_img, w_img = frame.data.shape[:2]
margin = 0.2 # 20% expansion
for person in people:
# 1. Expand Box
pb = person.box
bw = pb.x2 - pb.x1
bh = pb.y2 - pb.y1
# Expanded coordinates for Step 2
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))
# Use expanded crop for cigarette detection
# 2. Crop
person_crop = frame.data[ey1:ey2, ex1:ex2]
# If crop is too small, skip
if (
person_crop.size == 0
or person_crop.shape[0] < 10
@@ -100,10 +114,10 @@ class SmokingDetectionPipeline:
):
continue
# 3. Detect
cig_candidates = self.cigarette_model.predict(person_crop)
# Filter out extremely small candidates (noise)
# User requested check similar to person_crop check (< 10px)
# 4. Filter Noise
cig_candidates = [
c
for c in cig_candidates
@@ -111,25 +125,10 @@ class SmokingDetectionPipeline:
]
if cig_candidates:
# Store expansion offset for later coordinate translation
# We store (ex1, ey1) as metadata on the person detection or passed tuple?
# Easiest is to pass it in the tuple: (person, cigs, offset)
candidates.append((person, cig_candidates, (ex1, ey1)))
# Check detections (translate to global coordinates)
# 5. Global Coordinate Mapping (For Visualization)
for cig in cig_candidates:
# Clone or modify the cig detection to be global
# Assuming cig.box is relevant to person_crop
# Global X = person_crop_origin_x + local_x
# Global Y = person_crop_origin_y + local_y
# Update in place might affect other logic if it assumes local,
# but logic below (Expansion) seems to rely on cig.box being used on raw_person_crop
# which implies local coordinates.
# So we should create a COPY for visualization.
from pipeline.types import BoundingBox, Detection
global_cig = Detection(
box=BoundingBox(
x1=ex1 + cig.box.x1,
@@ -144,97 +143,47 @@ class SmokingDetectionPipeline:
)
frame.detections.append(global_cig)
if not candidates:
return # EXIT
return candidates
logger.debug(f"Frame {frame.frame_id}: Found {len(candidates)} candidates.")
# --- Step 3: Verification (Classifier) ---
# "Is that small thing actually a cigarette?"
def _verify_candidates(
self, frame: Frame, candidates: List[Tuple]
) -> Optional[Tuple]:
"""Runs verification on candidate crops."""
confirmed_hits = []
for person, cigs, (ex1, ey1) in candidates:
for cig in cigs:
# Crop just the object.
# Note: box coordinates are relative to the person crop, or global?
# Ideally, detections should be unified.
# If cigarette_model returns boxes relative to person_crop, we need to adjust them.
# For simplicity in this mock, let's assume models return GLOBAL coords or handle their own relative logic wrapping.
# However, usually crop inference returns local coords.
# Let's assume for now the cig model knows how to return global.
# OR we pass the crop to verifier.
# IMPORTANT: If cigarette model ran on a crop, its output box is relative to that crop.
# We need to map it back or just use the crop directly?
# Actually, we likely want to crop the cigarette from the person crop.
# Simplified: Just rely on the verifier taking the small patch.
# To be robust, let's assume the Detector returns a Detection with a convenient 'crop' method or we handle coordinate translation.
# But for this MVP, let's assume we pass the crop derived from the box.
# FIX: If cig_candidates came from person_crop, their boxes are relative to (0,0) of person_crop.
# We should probably run verification on the crop defined by cig.box within person_crop.
# Let's grab the actual image patch for verification
# We can'teasily re-crop from global frame without translation.
# Strategy: cig detection on person_crop is easiest.
# Re-cutting the crop:
# We need the SAME crop used for detection, or at least the same coordinate system.
# Since we used the expanded crop (person_crop created via ex1, ey1),
# cig.box is relative to that.
# We can just re-extract the crop from the frame using ex1, ey1 if we didn't save it.
# Or just use the global coordinates we know now.
# Let's derive the object crop from the expanded person crop logic.
# --- Expansion Logic ---
# Expand the box to give context
# Here we are expanding the CIGARETTE box within the PERSON crop?
# The user asked to expand the hitbox (of the human) stage two too.
# We did that in Step 2.
# Now we need to prepare the crop for VERIFICATION (Step 3).
# We take the cig box (relative to expanded person crop) and add context?
# 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) # Add 50% width padding
pad_y = int(cig_h * 0.5) # Add 50% height padding
pad_x = int(cig_w * 0.5)
pad_y = int(cig_h * 0.5)
# Coordinates relative to the expanded person crop
cx1 = max(0, cig.box.x1 - pad_x)
cy1 = max(0, cig.box.y1 - pad_y)
# We don't know the exact W/H of the expanded crop easily here without recomputing or checking boundaries
# But we can map to Global and crop from frame directly to be safe.
global_cx1 = ex1 + cx1
global_cy1 = ey1 + cy1
global_cx1 = ex1 + max(0, cig.box.x1 - pad_x)
global_cy1 = ey1 + max(0, cig.box.y1 - pad_y)
global_cx2 = ex1 + cig.box.x2 + pad_x
global_cy2 = ey1 + cig.box.y2 + pad_y
# Check bounds
h_img, w_img = frame.data.shape[:2]
final_x1 = max(0, global_cx1)
final_y1 = max(0, global_cy1)
final_x2 = min(w_img, global_cx2)
final_y2 = min(h_img, global_cy2)
# 2. Crop
obj_crop = frame.data[final_y1:final_y2, final_x1:final_x2]
if obj_crop.size == 0:
continue
# 3. Verify
score = self.verifier_model.verify(obj_crop)
if score > 0.8: # High confidence threshold
if score > 0.8: # High confidence
confirmed_hits.append((person, cig, score))
# Add verified detection to frame for visualization
from pipeline.types import BoundingBox, Detection
verified_det = Detection(
box=BoundingBox(
x1=ex1 + cig.box.x1,
@@ -249,23 +198,6 @@ class SmokingDetectionPipeline:
)
frame.detections.append(verified_det)
elif score > 0.4 and self.vlm_model:
# Ambiguous case: Send to VLM?
vlm_score = self.vlm_model.verify(
obj_crop
) # Or pass seq of frames?
if vlm_score > 0.8:
confirmed_hits.append((person, cig, vlm_score))
# --- Step 4: Action ---
if confirmed_hits:
best_hit = max(confirmed_hits, key=lambda x: x[2])
ticket = Ticket(
frame_id=frame.frame_id,
confidence=best_hit[2],
evidence=[frame.data], # Store frame reference
stage_confirmed="VLM"
if self.vlm_model
and best_hit[2] == 0.1 # Logic here is a bit shaky from mock
else "Verifier",
)
self.store.store(ticket)
return max(confirmed_hits, key=lambda x: x[2])
return None
pipeline/models.py
+13 -3
View File
@@ -1,6 +1,7 @@
import logging
from typing import Any, List
import torch
from ultralytics.models import YOLO
from pipeline.protocols import Model, Verifier
@@ -8,6 +9,13 @@ from pipeline.types import BoundingBox, Detection
logger = logging.getLogger(__name__)
# Auto-detect device
DEVICE = 0 if torch.cuda.is_available() else "cpu"
if DEVICE == 0:
logger.info(f"GPU detected! Using CUDA device: {torch.cuda.get_device_name(0)}")
else:
logger.info("No GPU detected. Running on CPU.")
class YoloPersonDetector(Model):
def __init__(self, model_path: str = "yolo11n.pt"): # Using small model for speed
@@ -16,7 +24,9 @@ class YoloPersonDetector(Model):
def predict(self, input_data: Any) -> List[Detection]:
# input_data: numpy array (H, W, C)
results = self.model(input_data, classes=[0], verbose=False) # 0 is person
results = self.model(
input_data, classes=[0], verbose=False, device=DEVICE
) # 0 is person
detections = []
for r in results:
@@ -42,7 +52,7 @@ class YoloCigaretteDetector(Model):
def predict(self, input_data: Any) -> List[Detection]:
# This model should be trained on cigarette class
# Assuming class 0 is cigarette for this custom model
results = self.model(input_data, verbose=False)
results = self.model(input_data, verbose=False, device=DEVICE)
detections = []
for r in results:
@@ -72,7 +82,7 @@ class YoloClassifier(Verifier):
if input_data.size == 0:
return 0.0
results = self.model(input_data, verbose=False)
results = self.model(input_data, verbose=False, device=DEVICE)
# Check if we got any results
if not results:
pipeline/outputs.py
+7 -3
View File
@@ -6,7 +6,7 @@ from typing import List, Optional
import cv2
import numpy as np
from pipeline.types import Detection, Frame, Ticket
from pipeline.types import Detection, Frame
logger = logging.getLogger(__name__)
@@ -19,10 +19,14 @@ class FileAlertStore:
def __init__(self, path: str):
self.path = path
self.file = open(path, "a") # Append mode
self.file = open(path, "a", encoding="utf-8") # Append mode
logger.info(f"Initialized FileAlertStore writing to {path}")
def store(self, ticket: Ticket):
def handle(self, frame: Frame):
if not frame.ticket:
return
ticket = frame.ticket
record = {
"type": "alert",
"frame_id": ticket.frame_id,
pipeline/protocols.py
+4 -4
View File
@@ -1,6 +1,6 @@
from typing import Any, Iterator, List, Protocol
from pipeline.types import Detection, Frame, Ticket
from pipeline.types import Detection, Frame
class Model(Protocol):
@@ -21,7 +21,7 @@ class StreamSource(Protocol):
...
class AlertStore(Protocol):
def store(self, ticket: Ticket):
"""Persist the ticket."""
class Sink(Protocol):
def handle(self, frame: Frame):
"""Process or persist a frame (and its ticket/detections)."""
...
pipeline/types.py
+12 -11
View File
@@ -1,5 +1,5 @@
from dataclasses import dataclass, field
from typing import Any, List
from typing import Any, List, Optional
import numpy as np
@@ -39,16 +39,6 @@ class Detection:
metadata: dict = field(default_factory=dict)
@dataclass
class Frame:
"""Represents a single frame in the video stream."""
timestamp: float
data: np.ndarray # numpy array (OpenCV image)
frame_id: int
detections: List[Detection] = field(default_factory=list)
@dataclass
class Ticket:
"""Represents a potential alert event."""
@@ -58,3 +48,14 @@ class Ticket:
evidence: List[Any] # List of images or crops
stage_confirmed: str
metadata: dict = field(default_factory=dict)
@dataclass
class Frame:
"""Represents a single frame in the video stream."""
timestamp: float
data: np.ndarray # numpy array (OpenCV image)
frame_id: int
detections: List[Detection] = field(default_factory=list)
ticket: Optional[Ticket] = None