bout right
This commit is contained in:
1 parent
34bee8bf31
commit
2fc7d1cbfa
12 files changed
+1305
-335
No files matched your search
@@ -3,3 +3,4 @@
|
|||||||
*.jsonl filter=lfs diff=lfs merge=lfs -text
|
*.jsonl filter=lfs diff=lfs merge=lfs -text
|
||||||
*.png filter=lfs diff=lfs merge=lfs -text
|
*.png filter=lfs diff=lfs merge=lfs -text
|
||||||
*.jpg filter=lfs diff=lfs merge=lfs -text
|
*.jpg filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.db filter=lfs diff=lfs merge=lfs -text
|
||||||
@@ -8,5 +8,3 @@ wheels/
|
|||||||
|
|
||||||
# Virtual environments
|
# Virtual environments
|
||||||
.venv
|
.venv
|
||||||
|
|
||||||
Group_project_*
|
|
||||||
Submodule
+1
Submodule SmokeDetectionApp added at f20e9a7c76.
@@ -0,0 +1,18 @@
|
|||||||
|
"""Database abstraction layer.
|
||||||
|
|
||||||
|
Provides a Repository protocol and concrete implementations for SQLite and MySQL.
|
||||||
|
The server doesn't care which one it talks to.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from db.repository import Repository
|
||||||
|
from db.sqlite import SQLiteRepository
|
||||||
|
|
||||||
|
__all__ = ["Repository", "SQLiteRepository"]
|
||||||
|
|
||||||
|
# Optional MySQL import (requires mysql-connector-python)
|
||||||
|
try:
|
||||||
|
from db.mysql import MySQLRepository
|
||||||
|
|
||||||
|
__all__.append("MySQLRepository")
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
+135
@@ -0,0 +1,135 @@
|
|||||||
|
"""MySQL implementation of the Repository."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
import mysql.connector # type: ignore[import-untyped]
|
||||||
|
|
||||||
|
from db.repository import DetectionRecord, User
|
||||||
|
|
||||||
|
|
||||||
|
class MySQLRepository:
|
||||||
|
"""MySQL-backed repository. Needs a running MySQL server + .env config."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
user: Optional[str] = None,
|
||||||
|
password: Optional[str] = None,
|
||||||
|
host: Optional[str] = None,
|
||||||
|
database: Optional[str] = None,
|
||||||
|
):
|
||||||
|
self.config = {
|
||||||
|
"user": user or os.getenv("DB_USER", "smoke_admin"),
|
||||||
|
"password": password or os.getenv("DB_PASSWORD", "123456"),
|
||||||
|
"host": host or os.getenv("DB_HOST", "localhost"),
|
||||||
|
"database": database or os.getenv("DB_NAME", "smoke_detection"),
|
||||||
|
}
|
||||||
|
|
||||||
|
def _connect(self):
|
||||||
|
return mysql.connector.connect(**self.config)
|
||||||
|
|
||||||
|
def initialize(self) -> None:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute(
|
||||||
|
"""CREATE TABLE IF NOT EXISTS users (
|
||||||
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
||||||
|
username VARCHAR(50) UNIQUE NOT NULL,
|
||||||
|
password VARCHAR(255) NOT NULL
|
||||||
|
)"""
|
||||||
|
)
|
||||||
|
cursor.execute(
|
||||||
|
"""CREATE TABLE IF NOT EXISTS detections (
|
||||||
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
||||||
|
user_id INT,
|
||||||
|
image_path VARCHAR(255) NOT NULL,
|
||||||
|
confidence FLOAT,
|
||||||
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
camera_source VARCHAR(50) DEFAULT 'Laptop Webcam',
|
||||||
|
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||||
|
)"""
|
||||||
|
)
|
||||||
|
cursor.execute(
|
||||||
|
"INSERT IGNORE INTO users (username, password) VALUES ('admin', '123456')"
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
# --- Users ---
|
||||||
|
|
||||||
|
def get_user_by_username(self, username: str) -> Optional[User]:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor(dictionary=True)
|
||||||
|
cursor.execute("SELECT * FROM users WHERE username = %s", (username,))
|
||||||
|
row = cursor.fetchone()
|
||||||
|
conn.close()
|
||||||
|
return User(**row) if row else None
|
||||||
|
|
||||||
|
def get_user_by_credentials(self, username: str, password: str) -> Optional[User]:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor(dictionary=True)
|
||||||
|
cursor.execute(
|
||||||
|
"SELECT * FROM users WHERE username = %s AND password = %s",
|
||||||
|
(username, password),
|
||||||
|
)
|
||||||
|
row = cursor.fetchone()
|
||||||
|
conn.close()
|
||||||
|
return User(**row) if row else None
|
||||||
|
|
||||||
|
def create_user(self, username: str, password: str) -> int:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute(
|
||||||
|
"INSERT INTO users (username, password) VALUES (%s, %s)",
|
||||||
|
(username, password),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
new_id = cursor.lastrowid
|
||||||
|
conn.close()
|
||||||
|
return new_id # type: ignore[return-value]
|
||||||
|
|
||||||
|
# --- Detections ---
|
||||||
|
|
||||||
|
def save_detection(
|
||||||
|
self,
|
||||||
|
user_id: int,
|
||||||
|
image_path: str,
|
||||||
|
confidence: float,
|
||||||
|
camera_source: str = "Laptop Webcam",
|
||||||
|
) -> int:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute(
|
||||||
|
"INSERT INTO detections (user_id, image_path, confidence, camera_source) VALUES (%s, %s, %s, %s)",
|
||||||
|
(user_id, image_path, confidence, camera_source),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
new_id = cursor.lastrowid
|
||||||
|
conn.close()
|
||||||
|
return new_id # type: ignore[return-value]
|
||||||
|
|
||||||
|
def get_detections(self, user_id: Optional[int] = None) -> list[DetectionRecord]:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor(dictionary=True)
|
||||||
|
if user_id is not None:
|
||||||
|
cursor.execute(
|
||||||
|
"SELECT * FROM detections WHERE user_id = %s ORDER BY timestamp DESC",
|
||||||
|
(user_id,),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
cursor.execute("SELECT * FROM detections ORDER BY timestamp DESC")
|
||||||
|
rows = cursor.fetchall()
|
||||||
|
conn.close()
|
||||||
|
return [DetectionRecord(**r) for r in rows]
|
||||||
|
|
||||||
|
def clear_detections(self, user_id: int) -> int:
|
||||||
|
conn = self._connect()
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("DELETE FROM detections WHERE user_id = %s", (user_id,))
|
||||||
|
conn.commit()
|
||||||
|
count = cursor.rowcount
|
||||||
|
conn.close()
|
||||||
|
return count
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
pass # MySQL uses per-call connections
|
||||||
@@ -0,0 +1,68 @@
|
|||||||
|
"""Repository protocol - the interface all DB backends must implement."""
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Optional, Protocol, runtime_checkable
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class User:
|
||||||
|
id: int
|
||||||
|
username: str
|
||||||
|
password: str # plaintext for now (legacy compat) - hash this later
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DetectionRecord:
|
||||||
|
id: int
|
||||||
|
user_id: int
|
||||||
|
image_path: str
|
||||||
|
confidence: float
|
||||||
|
timestamp: str
|
||||||
|
camera_source: str
|
||||||
|
|
||||||
|
|
||||||
|
@runtime_checkable
|
||||||
|
class Repository(Protocol):
|
||||||
|
"""Database operations for the smoking detection server."""
|
||||||
|
|
||||||
|
def initialize(self) -> None:
|
||||||
|
"""Create tables if they don't exist."""
|
||||||
|
...
|
||||||
|
|
||||||
|
# --- Users ---
|
||||||
|
|
||||||
|
def get_user_by_username(self, username: str) -> Optional[User]:
|
||||||
|
"""Find a user by username. Returns None if not found."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def get_user_by_credentials(self, username: str, password: str) -> Optional[User]:
|
||||||
|
"""Find a user by username + password. Returns None if invalid."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def create_user(self, username: str, password: str) -> int:
|
||||||
|
"""Create a new user. Returns the new user ID."""
|
||||||
|
...
|
||||||
|
|
||||||
|
# --- Detections ---
|
||||||
|
|
||||||
|
def save_detection(
|
||||||
|
self,
|
||||||
|
user_id: int,
|
||||||
|
image_path: str,
|
||||||
|
confidence: float,
|
||||||
|
camera_source: str = "Laptop Webcam",
|
||||||
|
) -> int:
|
||||||
|
"""Save a detection record. Returns the new record ID."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def get_detections(self, user_id: Optional[int] = None) -> list[DetectionRecord]:
|
||||||
|
"""Get detection history, optionally filtered by user. Newest first."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def clear_detections(self, user_id: int) -> int:
|
||||||
|
"""Delete all detections for a user. Returns number of rows deleted."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""Close the database connection."""
|
||||||
|
...
|
||||||
+121
@@ -0,0 +1,121 @@
|
|||||||
|
"""SQLite implementation of the Repository."""
|
||||||
|
|
||||||
|
import sqlite3
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from db.repository import DetectionRecord, User
|
||||||
|
|
||||||
|
# Schema that works with SQLite (no MySQL-specific syntax)
|
||||||
|
SCHEMA = """\
|
||||||
|
CREATE TABLE IF NOT EXISTS users (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
username TEXT UNIQUE NOT NULL,
|
||||||
|
password TEXT NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS detections (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
user_id INTEGER,
|
||||||
|
image_path TEXT NOT NULL,
|
||||||
|
confidence REAL,
|
||||||
|
timestamp TEXT DEFAULT (datetime('now', 'localtime')),
|
||||||
|
camera_source TEXT DEFAULT 'Laptop Webcam',
|
||||||
|
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||||
|
);
|
||||||
|
|
||||||
|
INSERT OR IGNORE INTO users (username, password) VALUES ('admin', '123456');
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class SQLiteRepository:
|
||||||
|
"""SQLite-backed repository. Zero setup, file-based."""
|
||||||
|
|
||||||
|
def __init__(self, db_path: str | Path = "data/smoke.db"):
|
||||||
|
self.db_path = Path(db_path)
|
||||||
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
self._conn: Optional[sqlite3.Connection] = None
|
||||||
|
|
||||||
|
def _get_conn(self) -> sqlite3.Connection:
|
||||||
|
if self._conn is None:
|
||||||
|
self._conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
|
||||||
|
self._conn.row_factory = sqlite3.Row
|
||||||
|
self._conn.execute("PRAGMA journal_mode=WAL")
|
||||||
|
self._conn.execute("PRAGMA foreign_keys=ON")
|
||||||
|
return self._conn
|
||||||
|
|
||||||
|
def initialize(self) -> None:
|
||||||
|
conn = self._get_conn()
|
||||||
|
conn.executescript(SCHEMA)
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
# --- Users ---
|
||||||
|
|
||||||
|
def get_user_by_username(self, username: str) -> Optional[User]:
|
||||||
|
row = (
|
||||||
|
self._get_conn()
|
||||||
|
.execute("SELECT * FROM users WHERE username = ?", (username,))
|
||||||
|
.fetchone()
|
||||||
|
)
|
||||||
|
return User(**dict(row)) if row else None
|
||||||
|
|
||||||
|
def get_user_by_credentials(self, username: str, password: str) -> Optional[User]:
|
||||||
|
row = (
|
||||||
|
self._get_conn()
|
||||||
|
.execute(
|
||||||
|
"SELECT * FROM users WHERE username = ? AND password = ?",
|
||||||
|
(username, password),
|
||||||
|
)
|
||||||
|
.fetchone()
|
||||||
|
)
|
||||||
|
return User(**dict(row)) if row else None
|
||||||
|
|
||||||
|
def create_user(self, username: str, password: str) -> int:
|
||||||
|
conn = self._get_conn()
|
||||||
|
cursor = conn.execute(
|
||||||
|
"INSERT INTO users (username, password) VALUES (?, ?)",
|
||||||
|
(username, password),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
return cursor.lastrowid # type: ignore[return-value]
|
||||||
|
|
||||||
|
# --- Detections ---
|
||||||
|
|
||||||
|
def save_detection(
|
||||||
|
self,
|
||||||
|
user_id: int,
|
||||||
|
image_path: str,
|
||||||
|
confidence: float,
|
||||||
|
camera_source: str = "Laptop Webcam",
|
||||||
|
) -> int:
|
||||||
|
conn = self._get_conn()
|
||||||
|
cursor = conn.execute(
|
||||||
|
"INSERT INTO detections (user_id, image_path, confidence, camera_source) VALUES (?, ?, ?, ?)",
|
||||||
|
(user_id, image_path, confidence, camera_source),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
return cursor.lastrowid # type: ignore[return-value]
|
||||||
|
|
||||||
|
def get_detections(self, user_id: Optional[int] = None) -> list[DetectionRecord]:
|
||||||
|
conn = self._get_conn()
|
||||||
|
if user_id is not None:
|
||||||
|
rows = conn.execute(
|
||||||
|
"SELECT * FROM detections WHERE user_id = ? ORDER BY timestamp DESC",
|
||||||
|
(user_id,),
|
||||||
|
).fetchall()
|
||||||
|
else:
|
||||||
|
rows = conn.execute(
|
||||||
|
"SELECT * FROM detections ORDER BY timestamp DESC"
|
||||||
|
).fetchall()
|
||||||
|
return [DetectionRecord(**dict(r)) for r in rows]
|
||||||
|
|
||||||
|
def clear_detections(self, user_id: int) -> int:
|
||||||
|
conn = self._get_conn()
|
||||||
|
cursor = conn.execute("DELETE FROM detections WHERE user_id = ?", (user_id,))
|
||||||
|
conn.commit()
|
||||||
|
return cursor.rowcount
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
if self._conn:
|
||||||
|
self._conn.close()
|
||||||
|
self._conn = None
|
||||||
@@ -1,6 +1,7 @@
|
|||||||
"""Sinks for pipeline output handling."""
|
"""Sinks for pipeline output handling."""
|
||||||
|
|
||||||
from pipeline.sinks.alerts import FileAlertStore
|
from pipeline.sinks.alerts import FileAlertStore
|
||||||
|
from pipeline.sinks.database import DatabaseSink
|
||||||
from pipeline.sinks.display import WindowDisplaySink
|
from pipeline.sinks.display import WindowDisplaySink
|
||||||
from pipeline.sinks.events import SmartEventSink
|
from pipeline.sinks.events import SmartEventSink
|
||||||
from pipeline.sinks.metrics import MetricsSink
|
from pipeline.sinks.metrics import MetricsSink
|
||||||
@@ -8,6 +9,7 @@ from pipeline.sinks.screenshots import ScreenshotSink
|
|||||||
from pipeline.sinks.video import VideoFileStore
|
from pipeline.sinks.video import VideoFileStore
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
|
"DatabaseSink",
|
||||||
"FileAlertStore",
|
"FileAlertStore",
|
||||||
"MetricsSink",
|
"MetricsSink",
|
||||||
"ScreenshotSink",
|
"ScreenshotSink",
|
||||||
|
|||||||
@@ -0,0 +1,74 @@
|
|||||||
|
"""Database sink - saves detection alerts to the database."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from db.repository import Repository
|
||||||
|
from pipeline.ops.drawing import draw_alert_banner, draw_detections
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class DatabaseSink:
|
||||||
|
"""
|
||||||
|
Pipeline sink that saves smoking alerts to a database.
|
||||||
|
|
||||||
|
On ticket (alert), saves an annotated screenshot and inserts a detection record.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
db: Repository,
|
||||||
|
user_id: int = 1,
|
||||||
|
image_dir: str = "static/alerts",
|
||||||
|
camera_source: str = "Pipeline",
|
||||||
|
):
|
||||||
|
self.db = db
|
||||||
|
self.user_id = user_id
|
||||||
|
self.camera_source = camera_source
|
||||||
|
self.image_dir = Path(image_dir)
|
||||||
|
self.image_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
self.count = 0
|
||||||
|
|
||||||
|
logger.info("DatabaseSink initialized (user_id=%d, dir=%s)", user_id, image_dir)
|
||||||
|
|
||||||
|
def handle(self, frame: Frame):
|
||||||
|
if not frame.ticket:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.count += 1
|
||||||
|
ticket = frame.ticket
|
||||||
|
|
||||||
|
# Save annotated image
|
||||||
|
filename = f"alert_{int(time.time())}_{self.count}.jpg"
|
||||||
|
filepath = self.image_dir / filename
|
||||||
|
|
||||||
|
img = frame.data.copy()
|
||||||
|
draw_detections(img, frame.detections)
|
||||||
|
draw_alert_banner(img, f"SMOKING {ticket.confidence:.0%}")
|
||||||
|
cv2.imwrite(str(filepath), img)
|
||||||
|
|
||||||
|
# Insert into database
|
||||||
|
web_path = f"/static/alerts/{filename}"
|
||||||
|
try:
|
||||||
|
self.db.save_detection(
|
||||||
|
user_id=self.user_id,
|
||||||
|
image_path=web_path,
|
||||||
|
confidence=ticket.confidence,
|
||||||
|
camera_source=self.camera_source,
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"Alert #%d saved to DB: %s (%.2f)",
|
||||||
|
self.count,
|
||||||
|
filename,
|
||||||
|
ticket.confidence,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to save alert to DB: %s", e)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
pass # DB connection is managed by the repository
|
||||||
@@ -5,11 +5,16 @@ description = "Add your description here"
|
|||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
requires-python = ">=3.13"
|
requires-python = ">=3.13"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"dotenv>=0.9.9",
|
||||||
|
"fastapi>=0.133.0",
|
||||||
|
"mysql-connector-python>=9.6.0",
|
||||||
"numpy>=2.4.0",
|
"numpy>=2.4.0",
|
||||||
"opencv-python>=4.11.0.86",
|
"opencv-python>=4.11.0.86",
|
||||||
|
"python-dotenv>=1.2.1",
|
||||||
"torch>=2.9.1",
|
"torch>=2.9.1",
|
||||||
"torchvision>=0.24.1",
|
"torchvision>=0.24.1",
|
||||||
"ultralytics>=8.3.248",
|
"ultralytics>=8.3.248",
|
||||||
|
"uvicorn>=0.41.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[tool.uv.index]]
|
[[tool.uv.index]]
|
||||||
|
|||||||
@@ -0,0 +1,308 @@
|
|||||||
|
"""
|
||||||
|
Smoking Detection Server
|
||||||
|
========================
|
||||||
|
|
||||||
|
FastAPI server that runs the full detection pipeline on a live webcam feed
|
||||||
|
and exposes REST endpoints for the frontend (auth, history, video stream).
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
uvicorn server:app --host 0.0.0.0 --port 8000
|
||||||
|
python server.py # dev mode
|
||||||
|
python server.py --staged # use stage-based pipeline
|
||||||
|
python server.py --source rtsp://... # RTSP camera
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
from fastapi import FastAPI, Form, Response
|
||||||
|
from fastapi.responses import StreamingResponse
|
||||||
|
from fastapi.staticfiles import StaticFiles
|
||||||
|
|
||||||
|
from db import SQLiteRepository
|
||||||
|
from db.repository import Repository
|
||||||
|
from pipeline.config import PipelineConfig
|
||||||
|
from pipeline.core import SmokingDetectionPipeline
|
||||||
|
from pipeline.models import YoloCigaretteDetector, YoloClassifier, YoloPersonDetector
|
||||||
|
from pipeline.protocols import Sink
|
||||||
|
from pipeline.sinks import DatabaseSink, FileAlertStore, MetricsSink
|
||||||
|
from pipeline.sources.stream import ThreadedStreamSource
|
||||||
|
from pipeline.staged_core import StageBasedPipeline
|
||||||
|
from pipeline.types import Frame
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Global state (set during lifespan)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
db: Repository = SQLiteRepository() # default: data/smoke.db
|
||||||
|
pipeline = None # set at startup
|
||||||
|
source: ThreadedStreamSource | None = None # set at startup
|
||||||
|
live_frame: Frame | None = None # latest frame for MJPEG / snapshot
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Sink that exposes frames to the web layer
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
class LiveFrameSink:
|
||||||
|
"""Captures the latest frame so API endpoints can stream it."""
|
||||||
|
|
||||||
|
def handle(self, frame: Frame):
|
||||||
|
global live_frame
|
||||||
|
live_frame = frame
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Pipeline runner (background thread)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
def _run_pipeline():
|
||||||
|
"""Blocking call: runs in a daemon thread."""
|
||||||
|
assert pipeline is not None and source is not None
|
||||||
|
try:
|
||||||
|
pipeline.process_stream(source)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
logger.exception("Pipeline crashed")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# CLI args (parsed before uvicorn when run as __main__)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
def parse_args():
|
||||||
|
p = argparse.ArgumentParser(description="Smoking Detection Server")
|
||||||
|
p.add_argument(
|
||||||
|
"--source", type=str, default="0", help="Webcam index or RTSP URL (default: 0)"
|
||||||
|
)
|
||||||
|
p.add_argument("--person-model", default="models/human-det-best.pt")
|
||||||
|
p.add_argument("--cig-model", default="models/cig-det-best.pt")
|
||||||
|
p.add_argument("--verify-model", default="models/cig-cls-best.pt")
|
||||||
|
p.add_argument("--staged", action="store_true", help="Use stage-based pipeline")
|
||||||
|
p.add_argument("--vlm", action="store_true")
|
||||||
|
p.add_argument("--vlm-key", default=None)
|
||||||
|
p.add_argument("--vlm-provider", default="openai", choices=["openai", "anthropic"])
|
||||||
|
p.add_argument("--vlm-model", default="gpt-4-vision-preview")
|
||||||
|
p.add_argument("--host", default="0.0.0.0")
|
||||||
|
p.add_argument("--port", type=int, default=8000)
|
||||||
|
p.add_argument("--debug", action="store_true")
|
||||||
|
return p.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
# Module-level defaults (overridden by __main__ block or env)
|
||||||
|
_args = argparse.Namespace(
|
||||||
|
source=os.environ.get("CAMERA_SOURCE", "0"),
|
||||||
|
person_model=os.environ.get("PERSON_MODEL", "models/human-det-best.pt"),
|
||||||
|
cig_model=os.environ.get("CIG_MODEL", "models/cig-det-best.pt"),
|
||||||
|
verify_model=os.environ.get("VERIFY_MODEL", "models/cig-cls-best.pt"),
|
||||||
|
staged=os.environ.get("STAGED", "").lower() in ("1", "true", "yes"),
|
||||||
|
vlm=os.environ.get("USE_VLM", "").lower() in ("1", "true", "yes"),
|
||||||
|
vlm_key=os.environ.get("VLM_API_KEY"),
|
||||||
|
vlm_provider=os.environ.get("VLM_PROVIDER", "openai"),
|
||||||
|
vlm_model=os.environ.get("VLM_MODEL", "gpt-4-vision-preview"),
|
||||||
|
debug=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Lifespan – startup & shutdown
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
@asynccontextmanager
|
||||||
|
async def app_lifespan(_app: FastAPI):
|
||||||
|
global pipeline, source # noqa: PLW0603
|
||||||
|
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.DEBUG if _args.debug else logging.INFO,
|
||||||
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||||
|
datefmt="%H:%M:%S",
|
||||||
|
)
|
||||||
|
logging.getLogger("ultralytics").setLevel(logging.WARNING)
|
||||||
|
|
||||||
|
# 1. Database
|
||||||
|
logger.info("Initializing database...")
|
||||||
|
db.initialize()
|
||||||
|
|
||||||
|
# 2. Models
|
||||||
|
logger.info("Loading models...")
|
||||||
|
person_model = YoloPersonDetector(_args.person_model)
|
||||||
|
cig_model = YoloCigaretteDetector(_args.cig_model)
|
||||||
|
verifier = YoloClassifier(_args.verify_model)
|
||||||
|
logger.info("Models loaded.")
|
||||||
|
|
||||||
|
# 3. Source (webcam / RTSP)
|
||||||
|
cam = int(_args.source) if _args.source.isdigit() else _args.source
|
||||||
|
source = ThreadedStreamSource(cam)
|
||||||
|
logger.info("Camera source ready: %s (%.1f FPS)", _args.source, source.fps)
|
||||||
|
|
||||||
|
# 4. Sinks
|
||||||
|
sinks: list[Sink] = [
|
||||||
|
LiveFrameSink(),
|
||||||
|
DatabaseSink(db, image_dir="static/alerts"),
|
||||||
|
MetricsSink(log_interval=200),
|
||||||
|
FileAlertStore("alerts.jsonl"),
|
||||||
|
]
|
||||||
|
|
||||||
|
# 5. Pipeline config
|
||||||
|
config = PipelineConfig(
|
||||||
|
idle_interval=0.5,
|
||||||
|
active_interval=0.0,
|
||||||
|
cooldown_frames=int(source.fps * 2),
|
||||||
|
)
|
||||||
|
|
||||||
|
# 6. Build pipeline
|
||||||
|
if _args.staged or _args.vlm:
|
||||||
|
vlm_key = _args.vlm_key or os.environ.get(
|
||||||
|
"OPENAI_API_KEY" if _args.vlm_provider == "openai" else "ANTHROPIC_API_KEY"
|
||||||
|
)
|
||||||
|
pipeline = StageBasedPipeline(
|
||||||
|
person_model=person_model,
|
||||||
|
cigarette_model=cig_model,
|
||||||
|
verifier_model=verifier,
|
||||||
|
vlm_api_key=vlm_key if _args.vlm else None,
|
||||||
|
vlm_provider=_args.vlm_provider,
|
||||||
|
vlm_model=_args.vlm_model,
|
||||||
|
use_vlm=_args.vlm,
|
||||||
|
sinks=sinks,
|
||||||
|
config=config,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
pipeline = SmokingDetectionPipeline(
|
||||||
|
person_model=person_model,
|
||||||
|
cigarette_model=cig_model,
|
||||||
|
verifier_model=verifier,
|
||||||
|
sinks=sinks,
|
||||||
|
config=config,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 7. Launch pipeline in background thread
|
||||||
|
t = threading.Thread(target=_run_pipeline, daemon=True, name="pipeline")
|
||||||
|
t.start()
|
||||||
|
logger.info("Pipeline running in background.")
|
||||||
|
|
||||||
|
yield # --- app serves requests ---
|
||||||
|
|
||||||
|
# Shutdown
|
||||||
|
logger.info("Shutting down...")
|
||||||
|
source.stop()
|
||||||
|
for s in sinks:
|
||||||
|
s.close()
|
||||||
|
db.close()
|
||||||
|
logger.info("Goodbye.")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# FastAPI app
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
app = FastAPI(title="Smoking Detection", lifespan=app_lifespan) # type: ignore
|
||||||
|
|
||||||
|
# Serve static assets (alert images, CSS/JS, etc.)
|
||||||
|
Path("static/alerts").mkdir(parents=True, exist_ok=True)
|
||||||
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
||||||
|
|
||||||
|
|
||||||
|
# -- Auth ------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/register")
|
||||||
|
def register(username: str = Form(...), password: str = Form(...)):
|
||||||
|
try:
|
||||||
|
if db.get_user_by_username(username):
|
||||||
|
return {"status": "fail", "message": "Username already taken"}
|
||||||
|
new_id = db.create_user(username, password)
|
||||||
|
return {"status": "success", "user_id": new_id, "message": "Account created"}
|
||||||
|
except Exception as e:
|
||||||
|
return {"status": "error", "message": str(e)}
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/login")
|
||||||
|
def login(username: str = Form(...), password: str = Form(...)):
|
||||||
|
try:
|
||||||
|
user = db.get_user_by_credentials(username, password)
|
||||||
|
if user:
|
||||||
|
return {"status": "success", "user_id": user.id, "username": user.username}
|
||||||
|
return {"status": "fail", "message": "Invalid username or password"}
|
||||||
|
except Exception as e:
|
||||||
|
return {"status": "error", "message": str(e)}
|
||||||
|
|
||||||
|
|
||||||
|
# -- Detection history -----------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/history")
|
||||||
|
def get_history(user_id: int):
|
||||||
|
try:
|
||||||
|
records = db.get_detections(user_id)
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"id": r.id,
|
||||||
|
"user_id": r.user_id,
|
||||||
|
"image_path": r.image_path,
|
||||||
|
"confidence": r.confidence,
|
||||||
|
"timestamp": r.timestamp,
|
||||||
|
"camera_source": r.camera_source,
|
||||||
|
}
|
||||||
|
for r in records
|
||||||
|
]
|
||||||
|
except Exception:
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/clear_history")
|
||||||
|
def clear_history(user_id: int = Form(...)):
|
||||||
|
try:
|
||||||
|
deleted = db.clear_detections(user_id)
|
||||||
|
logger.info("Cleared %d records for user %d", deleted, user_id)
|
||||||
|
return {"status": "success", "message": "History cleared"}
|
||||||
|
except Exception as e:
|
||||||
|
return {"status": "error", "message": str(e)}
|
||||||
|
|
||||||
|
|
||||||
|
# -- Video / Snapshot ------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/snapshot")
|
||||||
|
def snapshot():
|
||||||
|
"""Return the latest frame as a JPEG."""
|
||||||
|
if live_frame is not None:
|
||||||
|
_, jpg = cv2.imencode(".jpg", live_frame.data)
|
||||||
|
return Response(content=jpg.tobytes(), media_type="image/jpeg")
|
||||||
|
return Response(status_code=404)
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/video_feed")
|
||||||
|
def video_feed():
|
||||||
|
"""MJPEG stream – drop into an <img> tag."""
|
||||||
|
|
||||||
|
def generate():
|
||||||
|
while True:
|
||||||
|
if live_frame is not None:
|
||||||
|
from pipeline.ops.drawing import draw_detections
|
||||||
|
|
||||||
|
img = live_frame.data.copy()
|
||||||
|
draw_detections(img, live_frame.detections)
|
||||||
|
_, jpg = cv2.imencode(".jpg", img)
|
||||||
|
yield (
|
||||||
|
b"--frame\r\n"
|
||||||
|
b"Content-Type: image/jpeg\r\n\r\n" + jpg.tobytes() + b"\r\n"
|
||||||
|
)
|
||||||
|
time.sleep(0.033) # ~30 FPS cap
|
||||||
|
|
||||||
|
return StreamingResponse(
|
||||||
|
generate(), media_type="multipart/x-mixed-replace;boundary=frame"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Entry-point
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import uvicorn
|
||||||
|
|
||||||
|
_args = parse_args()
|
||||||
|
uvicorn.run(app, host=_args.host, port=_args.port)
|
||||||
Reference in new issue
Block a user