how do i signal people that it can take ints now
This commit is contained in:
1 parent
8f6c3969c0
commit
fa32c7b28c
6 files changed
+91
-23
No files matched your search
@@ -95,7 +95,10 @@ def main():
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Loading models... Person: {args.person_model}, Cig: {args.cig_model}, Verifier: {args.verify_model}"
|
"Loading models... Person: %s, Cig: %s, Verifier: %s",
|
||||||
|
args.person_model,
|
||||||
|
args.cig_model,
|
||||||
|
args.verify_model,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Initialize implementation-specific classes
|
# Initialize implementation-specific classes
|
||||||
@@ -104,7 +107,11 @@ def main():
|
|||||||
verifier = YoloClassifier(args.verify_model)
|
verifier = YoloClassifier(args.verify_model)
|
||||||
|
|
||||||
# Initialize Source first to get FPS
|
# Initialize Source first to get FPS
|
||||||
if args.rtsp:
|
# Detect webcam (int), RTSP (rtsp://), or file
|
||||||
|
if args.source.isdigit():
|
||||||
|
# Webcam index (0, 1, etc.)
|
||||||
|
source = ThreadedRTSPSource(int(args.source))
|
||||||
|
elif args.rtsp or args.source.startswith("rtsp://"):
|
||||||
# Use background thread for RTSP to prevent frame drop issues / lag
|
# Use background thread for RTSP to prevent frame drop issues / lag
|
||||||
source = ThreadedRTSPSource(args.source)
|
source = ThreadedRTSPSource(args.source)
|
||||||
else:
|
else:
|
||||||
|
|||||||
+4
-4
@@ -12,14 +12,14 @@ logger = logging.getLogger(__name__)
|
|||||||
# Auto-detect device
|
# Auto-detect device
|
||||||
DEVICE = 0 if torch.cuda.is_available() else "cpu"
|
DEVICE = 0 if torch.cuda.is_available() else "cpu"
|
||||||
if DEVICE == 0:
|
if DEVICE == 0:
|
||||||
logger.info(f"GPU detected! Using CUDA device: {torch.cuda.get_device_name(0)}")
|
logger.info("GPU detected! Using CUDA device: %s", torch.cuda.get_device_name(0))
|
||||||
else:
|
else:
|
||||||
logger.info("No GPU detected. Running on CPU.")
|
logger.info("No GPU detected. Running on CPU.")
|
||||||
|
|
||||||
|
|
||||||
class YoloPersonDetector(Model):
|
class YoloPersonDetector(Model):
|
||||||
def __init__(self, model_path: str = "yolo11n.pt"): # Using small model for speed
|
def __init__(self, model_path: str = "yolo11n.pt"): # Using small model for speed
|
||||||
logger.info(f"Loading Person Detector: {model_path}")
|
logger.info("Loading Person Detector: %s", model_path)
|
||||||
self.model = YOLO(model_path) # Load model
|
self.model = YOLO(model_path) # Load model
|
||||||
|
|
||||||
def predict(self, input_data: Any) -> List[Detection]:
|
def predict(self, input_data: Any) -> List[Detection]:
|
||||||
@@ -46,7 +46,7 @@ class YoloPersonDetector(Model):
|
|||||||
|
|
||||||
class YoloCigaretteDetector(Model):
|
class YoloCigaretteDetector(Model):
|
||||||
def __init__(self, model_path: str):
|
def __init__(self, model_path: str):
|
||||||
logger.info(f"Loading Cigarette Detector: {model_path}")
|
logger.info("Loading Cigarette Detector: %s", model_path)
|
||||||
self.model = YOLO(model_path)
|
self.model = YOLO(model_path)
|
||||||
|
|
||||||
def predict(self, input_data: Any) -> List[Detection]:
|
def predict(self, input_data: Any) -> List[Detection]:
|
||||||
@@ -74,7 +74,7 @@ class YoloCigaretteDetector(Model):
|
|||||||
|
|
||||||
class YoloClassifier(Verifier):
|
class YoloClassifier(Verifier):
|
||||||
def __init__(self, model_path: str):
|
def __init__(self, model_path: str):
|
||||||
logger.info(f"Loading Smoking Classifier: {model_path}")
|
logger.info("Loading Smoking Classifier: %s", model_path)
|
||||||
self.model = YOLO(model_path)
|
self.model = YOLO(model_path)
|
||||||
|
|
||||||
def verify(self, input_data: Any) -> float:
|
def verify(self, input_data: Any) -> float:
|
||||||
|
|||||||
+4
-4
@@ -46,7 +46,7 @@ class SmartEventSink(Sink):
|
|||||||
self.current_event_writer: Optional[cv2.VideoWriter] = None
|
self.current_event_writer: Optional[cv2.VideoWriter] = None
|
||||||
self.current_event_id = None
|
self.current_event_id = None
|
||||||
|
|
||||||
logger.info(f"Initialized SmartEventSink in {output_dir}")
|
logger.info("Initialized SmartEventSink in %s", output_dir)
|
||||||
|
|
||||||
def handle(self, frame: Frame):
|
def handle(self, frame: Frame):
|
||||||
# Always add to buffer if not recording to keep context
|
# Always add to buffer if not recording to keep context
|
||||||
@@ -146,7 +146,7 @@ class FileAlertStore:
|
|||||||
def __init__(self, path: str):
|
def __init__(self, path: str):
|
||||||
self.path = path
|
self.path = path
|
||||||
self.file = open(path, "a", encoding="utf-8") # Append mode
|
self.file = open(path, "a", encoding="utf-8") # Append mode
|
||||||
logger.info(f"Initialized FileAlertStore writing to {path}")
|
logger.info("Initialized FileAlertStore writing to %s", path)
|
||||||
|
|
||||||
def handle(self, frame: Frame):
|
def handle(self, frame: Frame):
|
||||||
if not frame.ticket:
|
if not frame.ticket:
|
||||||
@@ -201,7 +201,7 @@ class VideoFileStore:
|
|||||||
|
|
||||||
# Open metadata file
|
# Open metadata file
|
||||||
self.meta_file = open(self.metadata_path, "w")
|
self.meta_file = open(self.metadata_path, "w")
|
||||||
logger.info(f"Initialized Async VideoFileStore writing to {output_path}")
|
logger.info("Initialized Async VideoFileStore writing to %s", output_path)
|
||||||
self.thread.start()
|
self.thread.start()
|
||||||
|
|
||||||
def handle(self, frame: Frame):
|
def handle(self, frame: Frame):
|
||||||
@@ -224,7 +224,7 @@ class VideoFileStore:
|
|||||||
try:
|
try:
|
||||||
self._process_frame(frame)
|
self._process_frame(frame)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing video output: {e}")
|
logger.error("Error processing video output: %s", e)
|
||||||
finally:
|
finally:
|
||||||
self.queue.task_done()
|
self.queue.task_done()
|
||||||
|
|
||||||
|
|||||||
@@ -25,6 +25,10 @@ class StreamSource(Protocol):
|
|||||||
"""Yield frames from source."""
|
"""Yield frames from source."""
|
||||||
...
|
...
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
"""Stop the source and release resources."""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
class Sink(Protocol):
|
class Sink(Protocol):
|
||||||
def handle(self, frame: Frame):
|
def handle(self, frame: Frame):
|
||||||
|
|||||||
+17
-13
@@ -25,7 +25,7 @@ class FileStreamSource(StreamSource):
|
|||||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||||
cap.release()
|
cap.release()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Could not probe FPS from {path}: {e}")
|
logger.warning("Could not probe FPS from %s: %s", path, e)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def fps(self) -> float:
|
def fps(self) -> float:
|
||||||
@@ -40,7 +40,7 @@ class FileStreamSource(StreamSource):
|
|||||||
# Todo: handle image sequence
|
# Todo: handle image sequence
|
||||||
return
|
return
|
||||||
|
|
||||||
logger.info(f"Opening video file: {self.path}")
|
logger.info("Opening video file: %s", self.path)
|
||||||
cap = cv2.VideoCapture(self.path)
|
cap = cv2.VideoCapture(self.path)
|
||||||
frame_id = 0
|
frame_id = 0
|
||||||
while cap.isOpened():
|
while cap.isOpened():
|
||||||
@@ -53,31 +53,35 @@ class FileStreamSource(StreamSource):
|
|||||||
frame_id += 1
|
frame_id += 1
|
||||||
cap.release()
|
cap.release()
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
"""No-op for file source (stream ends naturally)."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
class ThreadedRTSPSource(StreamSource):
|
class ThreadedRTSPSource(StreamSource):
|
||||||
"""
|
"""
|
||||||
Reads from RTSP stream in a background thread to prevent buffer buildup
|
Reads from RTSP stream or webcam in a background thread to prevent buffer buildup
|
||||||
and allow main processing loop to run at its own pace without stalling the decoding.
|
and allow main processing loop to run at its own pace without stalling the decoding.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, rtsp_url: str, queue_size=5):
|
def __init__(self, source: str | int, queue_size: int = 5):
|
||||||
self.rtsp_url = rtsp_url
|
self.source = source
|
||||||
self.queue = Queue(maxsize=queue_size)
|
self.queue: Queue[Frame] = Queue(maxsize=queue_size)
|
||||||
self._fps = 30.0
|
self._fps = 30.0
|
||||||
|
|
||||||
# Probe FPS (with timeout safety ideally, but simple here)
|
# Probe FPS (with timeout safety ideally, but simple here)
|
||||||
try:
|
try:
|
||||||
cap = cv2.VideoCapture(self.rtsp_url)
|
cap = cv2.VideoCapture(self.source)
|
||||||
if cap.isOpened():
|
if cap.isOpened():
|
||||||
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
self._fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||||
cap.release()
|
cap.release()
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.warning("Could not probe FPS from RTSP source.")
|
logger.warning("Could not probe FPS from source.")
|
||||||
|
|
||||||
self.stopped = False
|
self.stopped = False
|
||||||
self.thread = threading.Thread(target=self._update, daemon=True)
|
self.thread = threading.Thread(target=self._update, daemon=True)
|
||||||
self._start()
|
self._start()
|
||||||
logger.info(f"Initialized ThreadedRTSPSource for {rtsp_url}")
|
logger.info("Initialized ThreadedRTSPSource for %s", source)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def fps(self) -> float:
|
def fps(self) -> float:
|
||||||
@@ -87,15 +91,15 @@ class ThreadedRTSPSource(StreamSource):
|
|||||||
self.thread.start()
|
self.thread.start()
|
||||||
|
|
||||||
def _update(self):
|
def _update(self):
|
||||||
logger.info(f"RTSP Thread started for {self.rtsp_url}")
|
logger.info("Reader thread started for %s", self.source)
|
||||||
cap = cv2.VideoCapture(self.rtsp_url)
|
cap = cv2.VideoCapture(self.source)
|
||||||
frame_id = 0
|
frame_id = 0
|
||||||
while not self.stopped:
|
while not self.stopped:
|
||||||
if not cap.isOpened():
|
if not cap.isOpened():
|
||||||
# Try reconnecting
|
# Try reconnecting
|
||||||
logger.warning("RTSP stream not open, retrying...")
|
logger.warning("Stream not open, retrying...")
|
||||||
time.sleep(1)
|
time.sleep(1)
|
||||||
cap = cv2.VideoCapture(self.rtsp_url)
|
cap = cv2.VideoCapture(self.source)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
ret, frame = cap.read()
|
ret, frame = cap.read()
|
||||||
|
|||||||
@@ -0,0 +1,53 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
"""Quick utility to check available webcams and their properties."""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
|
||||||
|
def check_webcams(max_index: int = 5):
|
||||||
|
"""Probe webcam indices 0 to max_index and print their info."""
|
||||||
|
print("Scanning for webcams...\n")
|
||||||
|
found = []
|
||||||
|
|
||||||
|
for i in range(max_index):
|
||||||
|
cap = cv2.VideoCapture(i)
|
||||||
|
if cap.isOpened():
|
||||||
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||||
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||||
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||||
|
backend = cap.getBackendName()
|
||||||
|
|
||||||
|
print(f"[{i}] ✓ FOUND")
|
||||||
|
print(f" Resolution: {width}x{height}")
|
||||||
|
print(f" FPS: {fps}")
|
||||||
|
print(f" Backend: {backend}")
|
||||||
|
|
||||||
|
# Try to read a frame to confirm it actually works
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if ret:
|
||||||
|
print(f" Frame read: OK ({frame.shape})")
|
||||||
|
else:
|
||||||
|
print(" Frame read: FAILED (camera may be in use)")
|
||||||
|
|
||||||
|
found.append(i)
|
||||||
|
print()
|
||||||
|
cap.release()
|
||||||
|
|
||||||
|
if not found:
|
||||||
|
print("No webcams found!")
|
||||||
|
else:
|
||||||
|
print(f"Found {len(found)} webcam(s): {found}")
|
||||||
|
print(f"\nUsage: uv run main.py --source {found[0]}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
if len(sys.argv) >= 2:
|
||||||
|
try:
|
||||||
|
m = int(sys.argv[1])
|
||||||
|
except ValueError:
|
||||||
|
print(f"Usage: {sys.argv[0]} [MAX_INDEX: 5]")
|
||||||
|
else:
|
||||||
|
check_webcams(m)
|
||||||
|
else:
|
||||||
|
check_webcams()
|
||||||
Reference in new issue
Block a user