comments typical of a gemini thinking
This commit is contained in:
1 parent
19f0b9f352
commit
b2f4b4f748
10 files changed
+131
-46
No files matched your search
+79
-22
@@ -7,6 +7,7 @@ from pipeline.types import Frame, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConsoleStore:
|
||||
def store(self, ticket: Ticket):
|
||||
logger.info(
|
||||
@@ -74,7 +75,22 @@ class SmokingDetectionPipeline:
|
||||
candidates = []
|
||||
for person in people:
|
||||
# ROI: Region of Interest (e.g., upper body or hands)
|
||||
person_crop = person.box.crop(frame.data)
|
||||
# EXPANSTION: Expand the person box to include extended arms/surroundings
|
||||
h_img, w_img = frame.data.shape[:2]
|
||||
margin = 0.2 # 20% expansion
|
||||
|
||||
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
|
||||
person_crop = frame.data[ey1:ey2, ex1:ex2]
|
||||
|
||||
# If crop is too small, skip
|
||||
if (
|
||||
@@ -85,15 +101,27 @@ class SmokingDetectionPipeline:
|
||||
continue
|
||||
|
||||
cig_candidates = self.cigarette_model.predict(person_crop)
|
||||
|
||||
# Filter out extremely small candidates (noise)
|
||||
# User requested check similar to person_crop check (< 10px)
|
||||
cig_candidates = [
|
||||
c
|
||||
for c in cig_candidates
|
||||
if (c.box.x2 - c.box.x1) >= 10 and (c.box.y2 - c.box.y1) >= 10
|
||||
]
|
||||
|
||||
if cig_candidates:
|
||||
candidates.append((person, 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)
|
||||
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.box.x1 + local_x
|
||||
# Global Y = person.box.y1 + local_y
|
||||
# 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
|
||||
@@ -104,10 +132,10 @@ class SmokingDetectionPipeline:
|
||||
|
||||
global_cig = Detection(
|
||||
box=BoundingBox(
|
||||
x1=person.box.x1 + cig.box.x1,
|
||||
y1=person.box.y1 + cig.box.y1,
|
||||
x2=person.box.x1 + cig.box.x2,
|
||||
y2=person.box.y1 + cig.box.y2,
|
||||
x1=ex1 + cig.box.x1,
|
||||
y1=ey1 + cig.box.y1,
|
||||
x2=ex1 + cig.box.x2,
|
||||
y2=ey1 + cig.box.y2,
|
||||
conf=cig.box.conf,
|
||||
),
|
||||
score=cig.score,
|
||||
@@ -124,7 +152,7 @@ class SmokingDetectionPipeline:
|
||||
# --- Step 3: Verification (Classifier) ---
|
||||
# "Is that small thing actually a cigarette?"
|
||||
confirmed_hits = []
|
||||
for person, cigs in candidates:
|
||||
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?
|
||||
@@ -152,20 +180,49 @@ class SmokingDetectionPipeline:
|
||||
# Strategy: cig detection on person_crop is easiest.
|
||||
|
||||
# Re-cutting the crop:
|
||||
raw_person_crop = person.box.crop(frame.data)
|
||||
# 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
|
||||
h, w = raw_person_crop.shape[:2]
|
||||
pad_x = int((cig.box.x2 - cig.box.x1) * 0.5) # Add 50% width padding
|
||||
pad_y = int((cig.box.y2 - cig.box.y1) * 0.5) # Add 50% height padding
|
||||
# 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.
|
||||
|
||||
x1 = max(0, cig.box.x1 - pad_x)
|
||||
y1 = max(0, cig.box.y1 - pad_y)
|
||||
x2 = min(w, cig.box.x2 + pad_x)
|
||||
y2 = min(h, cig.box.y2 + pad_y)
|
||||
# Now we need to prepare the crop for VERIFICATION (Step 3).
|
||||
# We take the cig box (relative to expanded person crop) and add context?
|
||||
|
||||
obj_crop = raw_person_crop[y1:y2, x1:x2]
|
||||
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
|
||||
|
||||
# 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_cx2 = ex1 + cig.box.x2 + pad_x
|
||||
global_cy2 = ey1 + cig.box.y2 + pad_y
|
||||
|
||||
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)
|
||||
|
||||
obj_crop = frame.data[final_y1:final_y2, final_x1:final_x2]
|
||||
|
||||
if obj_crop.size == 0:
|
||||
continue
|
||||
@@ -180,10 +237,10 @@ class SmokingDetectionPipeline:
|
||||
|
||||
verified_det = Detection(
|
||||
box=BoundingBox(
|
||||
x1=person.box.x1 + cig.box.x1,
|
||||
y1=person.box.y1 + cig.box.y1,
|
||||
x2=person.box.x1 + cig.box.x2,
|
||||
y2=person.box.y1 + cig.box.y2,
|
||||
x1=ex1 + cig.box.x1,
|
||||
y1=ey1 + cig.box.y1,
|
||||
x2=ex1 + cig.box.x2,
|
||||
y2=ey1 + cig.box.y2,
|
||||
conf=score,
|
||||
),
|
||||
score=score,
|
||||
|
||||
+33
-1
@@ -6,11 +6,43 @@ from typing import List, Optional
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from pipeline.types import Detection, Frame
|
||||
from pipeline.types import Detection, Frame, Ticket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileAlertStore:
|
||||
"""
|
||||
Persists actionable alerts (Tickets) to a JSONL file.
|
||||
Use this for integrating with backend systems or audit logs.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = path
|
||||
self.file = open(path, "a") # Append mode
|
||||
logger.info(f"Initialized FileAlertStore writing to {path}")
|
||||
|
||||
def store(self, ticket: Ticket):
|
||||
record = {
|
||||
"type": "alert",
|
||||
"frame_id": ticket.frame_id,
|
||||
"confidence": ticket.confidence,
|
||||
"stage": ticket.stage_confirmed,
|
||||
"metadata": ticket.metadata,
|
||||
}
|
||||
# Note: We skip serializing 'evidence' (raw arrays) to JSON
|
||||
# In a real app, you'd save the image to disk and link the path here.
|
||||
|
||||
self.file.write(json.dumps(record) + "\n")
|
||||
self.file.flush()
|
||||
logger.info(
|
||||
f"[ALERT SAVED] {ticket.confidence:.2f} confirmed at {ticket.stage_confirmed}"
|
||||
)
|
||||
|
||||
def close(self):
|
||||
self.file.close()
|
||||
|
||||
|
||||
class VideoFileStore:
|
||||
"""
|
||||
Writes processed frames to a video file with visualization overlays.
|
||||
|
||||
Reference in new issue
Block a user