diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..e07ec97 --- /dev/null +++ b/.env.example @@ -0,0 +1,14 @@ +# Скопируйте в .env и заполните. Файл .env не попадает в git. + +# Telegram: токен от @BotFather и ID чата (узнать можно у @userinfobot) +ECOWATCH_TG_TOKEN= +ECOWATCH_TG_CHAT_ID= + +# Координаты камеры — для карты инцидентов (необязательно) +ECOWATCH_CAMERA_LAT= +ECOWATCH_CAMERA_LON= + +# Необязательно: путь к модели, папка данных, устройство (cpu / 0 для GPU) +# ECOWATCH_MODEL=models/best.pt +# ECOWATCH_DATA_DIR=data +# ECOWATCH_DEVICE=0 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..2f9343a --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,19 @@ +name: tests + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + # Тесты не требуют ultralytics/torch — ставим только лёгкие зависимости. + - run: pip install numpy opencv-python-headless Pillow requests -r requirements-dev.txt + - run: ruff check . + - run: ruff format --check . + - run: pytest -q diff --git a/.gitignore b/.gitignore index 28c577c..fe555bb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,26 +1,27 @@ .venv/ +venv/ .idea/ .vscode/ __pycache__/ *.pyc +.pytest_cache/ +.ruff_cache/ +# секреты +.env + +# данные приложения: журнал, снимки, клипы, загрузки +data/ +recordings/ +detection_history.json + +# датасеты и результаты обучения dataset_detect/ dataset_retrain/ fire_smoke_images/ +runs/ runs_retrain/ - shop project/ *.pt !models/best.pt - -recordings/ -detection_history.json - -1111.py -1111_backup*.py -check_classes.py -check_imports.py -download_extra_data.py -retrain_model.py -train_detect.py diff --git a/.streamlit/config.toml b/.streamlit/config.toml new file mode 100644 index 0000000..6fbc758 --- /dev/null +++ b/.streamlit/config.toml @@ -0,0 +1,15 @@ +[theme] +base = "dark" +primaryColor = "#ef4444" +backgroundColor = "#0f1117" +secondaryBackgroundColor = "#1a1d26" +textColor = "#e2e4eb" + +[client] +toolbarMode = "minimal" + +[browser] +gatherUsageStats = false + +[server] +maxUploadSize = 1024 diff --git a/README.md b/README.md index f75b483..c66fa42 100644 --- a/README.md +++ b/README.md @@ -1,64 +1,132 @@ -# EcoWatch AI +# 🛡️ EcoWatch AI -Система обнаружения экологических угроз на основе YOLOv8. +Система обнаружения пожаров и дыма на видео в реальном времени на базе YOLOv8. -Анализирует видеопоток в реальном времени и определяет: -- 🔥 Пожар -- 💨 Дым -- 🗑️ Сжигание мусора +Работает с веб-камерой, IP-камерой (RTSP/HTTP), видеофайлом или захватом экрана. +Подтверждённые угрозы записываются в журнал со снимком и видеоклипом, а уведомления уходят в Telegram. ## Возможности -- Мониторинг с веб-камеры, IP-камеры, экрана или загруженного файла -- Уведомления в Telegram с фото -- Аналитика и графики -- Журнал обнаружений с экспортом в CSV -- Запись видео при обнаружении угрозы +- **Без ложных «миганий».** Вероятности сглаживаются по времени, а инцидент + начинается только после нескольких подряд опасных кадров и заканчивается после серии чистых. + Если дым перерастает в пожар, это фиксируется в том же инциденте. +- **Инциденты, а не спам.** Одна угроза — одна запись в журнале (начало, конец, пиковая + вероятность, снимок, клип) и одно уведомление в Telegram с настраиваемой паузой. +- **Фоновый мониторинг.** Видео обрабатывается в отдельном потоке: интерфейс не подвисает, + а мониторинг продолжается при переключении вкладок и перезагрузке страницы. +- **Клипы с предзаписью.** В клип попадают несколько секунд *до* тревоги. +- **Переподключение к камере** при обрыве связи. +- **Консольный режим** для круглосуточной работы на сервере без браузера. +- Аналитика по периодам, тепловая карта активности, карта камер, экспорт журнала в CSV. +- Поддерживаются модели classify (как текущая `models/best.pt`) и detect (с рамками). ## Установка +Нужен Python 3.10 или новее. + ```bash +python -m venv .venv +.venv\Scripts\activate # Windows +source .venv/bin/activate # Linux/macOS pip install -r requirements.txt ``` +### Настройка Telegram + +Скопируйте `.env.example` в `.env` и заполните: + +```ini +ECOWATCH_TG_TOKEN=123456:ABC... # токен от @BotFather +ECOWATCH_TG_CHAT_ID=123456789 # ID чата, можно узнать у @userinfobot +``` + +Токен никогда не хранится в коде: файл `.env` исключён из git. + ## Запуск +### Веб-интерфейс + ```bash streamlit run main.py ``` -Или через батник: +Или двойной клик по `run_app.bat` в Windows. Откроется http://localhost:8501. +### Консольный режим + +```bash +python -m ecowatch 0 # веб-камера №0 +python -m ecowatch rtsp://user:pass@192.168.1.10/live --record +python -m ecowatch video.mp4 --show # с окном предпросмотра +python -m ecowatch screen # захват экрана +python -m ecowatch report # отчёт за сегодня в Telegram +python -m ecowatch --help # все параметры ``` -run_app.bat -``` + +Ежедневный отчёт удобно отправлять по расписанию: через планировщик заданий Windows или cron. + +## Настройка чувствительности + +| Параметр | По умолчанию | Смысл | +|---|---|---| +| Порог вероятности | 0.65 | минимальная сглаженная вероятность опасного класса | +| Отрыв от 2-го класса | 0.15 | защита от неоднозначных кадров | +| Сглаживание | 0.35 | вес нового кадра: меньше — стабильнее, но медленнее реакция | +| Кадров для подтверждения | 4 | сколько опасных кадров подряд нужно для тревоги | +| Кадров для отбоя | 15 | сколько чистых кадров подряд завершают инцидент | + +В веб-интерфейсе пороги меняются на лету, без перезапуска мониторинга. ## Обучение модели -Подготовьте датасет в папке `dataset_detect/` со структурой: +**Классификация** (как у текущей модели): датасет из папок по классам. ``` dataset_detect/ -├── train/ -│ ├── fire/ -│ ├── normal/ -│ └── smoke/ -└── valid/ - ├── fire/ - ├── normal/ - └── smoke/ +├── train/ {fire, normal, smoke}/ +└── val/ {fire, normal, smoke}/ +``` + +```bash +python train.py --data dataset_detect ``` -Запустите: +**Детекция с рамками** (рекомендуется: точнее и показывает, *где* огонь): датасет в формате YOLO. ```bash -python train.py +python train.py --data fire_smoke/data.yaml --model yolov8s.pt --imgsz 640 +``` + +Лучшие веса автоматически копируются в `models/best.pt`, а прежняя модель сохраняется +рядом как `best_<дата>.pt`. + +## Структура проекта + +``` +main.py точка входа Streamlit +train.py обучение модели +ecowatch/ +├── config.py настройки и пороги (.env) +├── detector.py обёртка YOLO: classify и detect +├── tracker.py сглаживание и гистерезис → инциденты +├── monitor.py фоновый цикл мониторинга +├── sources.py камеры, потоки, файлы, экран +├── storage.py журнал инцидентов (SQLite) +├── notifier.py Telegram (в фоне) +├── recorder.py клипы с предзаписью +├── overlay.py отрисовка поверх видео +├── ui.py веб-интерфейс +└── cli.py консольный режим +tests/ тесты (pytest) ``` -## Стек +Данные хранятся в `data/`: база `ecowatch.db`, а также папки `snapshots/` и `clips/`. +Журнал старой версии (`detection_history.json`) импортируется автоматически при первом запуске. -- Python 3.11 -- Streamlit -- YOLOv8 (ultralytics) -- OpenCV -- Plotly +## Разработка + +```bash +pip install -r requirements-dev.txt +pytest +ruff check . && ruff format --check . +``` diff --git a/ecowatch/__init__.py b/ecowatch/__init__.py new file mode 100644 index 0000000..9ee53b8 --- /dev/null +++ b/ecowatch/__init__.py @@ -0,0 +1,3 @@ +"""EcoWatch AI — обнаружение пожаров и дыма на видео с помощью YOLO.""" + +__version__ = "4.0.0" diff --git a/ecowatch/__main__.py b/ecowatch/__main__.py new file mode 100644 index 0000000..dd8a8c9 --- /dev/null +++ b/ecowatch/__main__.py @@ -0,0 +1,5 @@ +import sys + +from .cli import main + +sys.exit(main()) diff --git a/ecowatch/cli.py b/ecowatch/cli.py new file mode 100644 index 0000000..c4caee8 --- /dev/null +++ b/ecowatch/cli.py @@ -0,0 +1,114 @@ +"""Консольный режим без браузера — для круглосуточной работы на сервере. + +Примеры: + python -m ecowatch 0 # веб-камера 0 + python -m ecowatch rtsp://user:pass@ip/stream --record + python -m ecowatch video.mp4 --show # с окном предпросмотра + python -m ecowatch report # отправить отчёт за сегодня в Telegram +""" + +from __future__ import annotations + +import argparse +import logging +import signal +import sys +import time + +from .config import DetectionParams, Settings +from .notifier import AsyncSender, TelegramNotifier, format_daily_report +from .sources import parse_source +from .storage import EventStore + + +def build_parser() -> argparse.ArgumentParser: + d = DetectionParams() + p = argparse.ArgumentParser(prog="ecowatch", description="EcoWatch AI — мониторинг пожаров и дыма") + p.add_argument("source", help="0 (камера), rtsp://…, путь к видео, screen[:N] или 'report'") + p.add_argument("--model", help="путь к весам YOLO (по умолчанию models/best.pt)") + p.add_argument("--conf", type=float, default=d.conf_threshold, help="порог вероятности") + p.add_argument("--margin", type=float, default=d.margin, help="отрыв от второго класса") + p.add_argument("--confirm", type=int, default=d.confirm_frames, help="кадров для подтверждения") + p.add_argument("--clear", type=int, default=d.clear_frames, help="кадров для отбоя") + p.add_argument("--record", action="store_true", help="записывать клипы инцидентов") + p.add_argument("--no-telegram", action="store_true", help="не отправлять уведомления") + p.add_argument("--cooldown", type=float, default=60, help="пауза между тревогами, с") + p.add_argument("--show", action="store_true", help="показывать окно с видео") + p.add_argument("-v", "--verbose", action="store_true") + return p + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)-7s %(message)s", + datefmt="%H:%M:%S", + ) + settings = Settings.from_env() + if args.model: + from pathlib import Path + + settings.model_path = Path(args.model) + settings.ensure_dirs() + store = EventStore(settings.db_path) + notifier = TelegramNotifier(settings.telegram_token, settings.telegram_chat_id) + + if args.source == "report": + if not notifier.enabled: + logging.error("Telegram не настроен (ECOWATCH_TG_TOKEN / ECOWATCH_TG_CHAT_ID)") + return 1 + return 0 if notifier.send_message(format_daily_report(store.day_summary())) else 1 + + from .detector import Detector + from .monitor import Monitor, MonitorOptions + + detector = Detector(settings.model_path, settings.device) + logging.info("Модель %s (%s): %s", settings.model_path.name, detector.task, ", ".join(detector.names)) + options = MonitorOptions( + source=parse_source(args.source), + params=DetectionParams( + conf_threshold=args.conf, margin=args.margin, confirm_frames=args.confirm, clear_frames=args.clear + ), + notify=not args.no_telegram and notifier.enabled, + notify_cooldown=args.cooldown, + record=args.record, + lat=settings.camera_lat, + lon=settings.camera_lon, + ) + if not args.no_telegram and not notifier.enabled: + logging.warning("Telegram не настроен — уведомления отключены") + sender = AsyncSender() + monitor = Monitor(detector, store, settings, options, notifier, sender) + + signal.signal(signal.SIGINT, lambda *_: monitor.stop(wait=False)) + if hasattr(signal, "SIGTERM"): + signal.signal(signal.SIGTERM, lambda *_: monitor.stop(wait=False)) + + monitor.start() + logging.info("Мониторинг запущен: %s (Ctrl+C — остановить)", options.source.label) + last_log = 0.0 + while monitor.alive: + s = monitor.state() + if args.show and s.frame is not None: + import cv2 + + cv2.imshow("EcoWatch AI", s.frame) + if cv2.waitKey(1) & 0xFF in (ord("q"), 27): + monitor.stop(wait=False) + else: + time.sleep(0.2) + if time.monotonic() - last_log > 30: + last_log = time.monotonic() + logging.info("кадров: %d · %.1f FPS · инцидентов: %d", s.frames, s.fps, s.incidents) + monitor.stop() + sender.join() + s = monitor.state() + if s.error: + logging.error("Ошибка: %s", s.error) + logging.info("Готово. Кадров: %d, инцидентов: %d", s.frames, s.incidents) + return 1 if s.error else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ecowatch/config.py b/ecowatch/config.py new file mode 100644 index 0000000..06f0164 --- /dev/null +++ b/ecowatch/config.py @@ -0,0 +1,94 @@ +"""Настройки приложения: пути, секреты и пороги обнаружения. + +Секреты (токен Telegram) читаются только из переменных окружения или файла +`.env` в корне проекта и никогда не хранятся в коде. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path + +ROOT_DIR = Path(__file__).resolve().parent.parent + + +def load_dotenv(path: Path) -> None: + """Минимальный парсер .env: KEY=VALUE, комментарии через #. Не перетирает окружение.""" + if not path.is_file(): + return + for raw in path.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip().removeprefix("export ").strip() + value = value.strip().strip('"').strip("'") + os.environ.setdefault(key, value) + + +def _float_or_none(value: str | None) -> float | None: + try: + return float(value) if value not in (None, "") else None + except ValueError: + return None + + +@dataclass +class DetectionParams: + """Пороги принятия решения. Подобраны под classify-модель из models/best.pt.""" + + conf_threshold: float = 0.65 # минимальная сглаженная вероятность опасного класса + margin: float = 0.15 # отрыв лидера от второго класса + box_conf: float = 0.35 # мин. уверенность рамки (только для detect-моделей) + ema_alpha: float = 0.35 # вес нового кадра при сглаживании вероятностей + confirm_frames: int = 4 # подряд «опасных» кадров для начала инцидента + clear_frames: int = 15 # подряд «чистых» кадров для завершения инцидента + + +@dataclass +class Settings: + model_path: Path = field(default_factory=lambda: ROOT_DIR / "models" / "best.pt") + data_dir: Path = field(default_factory=lambda: ROOT_DIR / "data") + telegram_token: str | None = None + telegram_chat_id: str | None = None + camera_lat: float | None = None + camera_lon: float | None = None + device: str | None = None # "cpu", "0" (GPU) и т.п.; None — автоматически + + @property + def db_path(self) -> Path: + return self.data_dir / "ecowatch.db" + + @property + def snapshots_dir(self) -> Path: + return self.data_dir / "snapshots" + + @property + def clips_dir(self) -> Path: + return self.data_dir / "clips" + + @property + def telegram_configured(self) -> bool: + return bool(self.telegram_token and self.telegram_chat_id) + + @classmethod + def from_env(cls, dotenv: Path | None = ROOT_DIR / ".env") -> Settings: + if dotenv is not None: + load_dotenv(dotenv) + env = os.environ + s = cls() + if env.get("ECOWATCH_MODEL"): + s.model_path = Path(env["ECOWATCH_MODEL"]).expanduser() + if env.get("ECOWATCH_DATA_DIR"): + s.data_dir = Path(env["ECOWATCH_DATA_DIR"]).expanduser() + s.telegram_token = env.get("ECOWATCH_TG_TOKEN") or None + s.telegram_chat_id = env.get("ECOWATCH_TG_CHAT_ID") or None + s.camera_lat = _float_or_none(env.get("ECOWATCH_CAMERA_LAT")) + s.camera_lon = _float_or_none(env.get("ECOWATCH_CAMERA_LON")) + s.device = env.get("ECOWATCH_DEVICE") or None + return s + + def ensure_dirs(self) -> None: + for d in (self.data_dir, self.snapshots_dir, self.clips_dir): + d.mkdir(parents=True, exist_ok=True) diff --git a/ecowatch/detector.py b/ecowatch/detector.py new file mode 100644 index 0000000..229bcea --- /dev/null +++ b/ecowatch/detector.py @@ -0,0 +1,97 @@ +"""Обёртка над YOLO: единый результат для classify- и detect-моделей.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + +from .hazards import get_hazard, normalize + + +@dataclass(frozen=True) +class Box: + x1: int + y1: int + x2: int + y2: int + conf: float + label: str + + +@dataclass +class FrameResult: + """Вероятности классов для одного кадра (+ рамки, если модель detect).""" + + scores: dict[str, float] + boxes: list[Box] = field(default_factory=list) + + def top(self) -> tuple[str, float]: + if not self.scores: + return "normal", 0.0 + label = max(self.scores, key=lambda k: (self.scores[k], get_hazard(k).severity)) + return label, self.scores[label] + + +def _to_numpy(x: Any) -> np.ndarray: + if hasattr(x, "cpu"): + x = x.cpu() + if hasattr(x, "numpy"): + x = x.numpy() + return np.asarray(x) + + +def parse_result(result: Any, box_conf: float = 0.35) -> FrameResult: + """Превращает `ultralytics.engine.results.Results` в `FrameResult`.""" + names = {int(k): normalize(v) for k, v in dict(result.names).items()} + + probs = getattr(result, "probs", None) + if probs is not None: + raw = _to_numpy(probs.data).astype(float).ravel() + return FrameResult({names[i]: float(raw[i]) for i in names if i < len(raw)}) + + boxes_obj = getattr(result, "boxes", None) + scores = {name: 0.0 for name in names.values()} + boxes: list[Box] = [] + if boxes_obj is not None and len(boxes_obj) > 0: + confs = _to_numpy(boxes_obj.conf).ravel() + clss = _to_numpy(boxes_obj.cls).ravel().astype(int) + xyxy = _to_numpy(boxes_obj.xyxy).reshape(-1, 4) + for conf, cls_id, (x1, y1, x2, y2) in zip(confs, clss, xyxy, strict=False): + if conf < box_conf: + continue + label = names.get(int(cls_id), f"class_{cls_id}") + boxes.append(Box(int(x1), int(y1), int(x2), int(y2), float(conf), label)) + scores[label] = max(scores.get(label, 0.0), float(conf)) + # У detect-модели нет класса «норма» — выводим его как отсутствие угрозы. + danger_max = max((v for k, v in scores.items() if get_hazard(k).danger), default=0.0) + scores["normal"] = max(scores.get("normal", 0.0), 1.0 - danger_max) + return FrameResult(scores, boxes) + + +class Detector: + def __init__(self, model_path: str | Path, device: str | None = None) -> None: + model_path = Path(model_path) + if not model_path.is_file(): + raise FileNotFoundError(f"Модель не найдена: {model_path}") + from ultralytics import YOLO # тяжёлый импорт — только когда реально нужен + + self.path = model_path + self.model = YOLO(str(model_path)) + self.task: str = self.model.task + self.names: list[str] = [normalize(n) for n in self.model.names.values()] + self.device = device + self._lock = threading.Lock() + + def predict(self, frame: np.ndarray, box_conf: float = 0.35) -> FrameResult: + with self._lock: + kwargs: dict[str, Any] = {"verbose": False} + if self.device: + kwargs["device"] = self.device + if self.task == "detect": + kwargs["conf"] = box_conf + result = self.model(frame, **kwargs)[0] + return parse_result(result, box_conf) diff --git a/ecowatch/hazards.py b/ecowatch/hazards.py new file mode 100644 index 0000000..69bb7ae --- /dev/null +++ b/ecowatch/hazards.py @@ -0,0 +1,52 @@ +"""Каталог классов, которые умеет различать модель.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Hazard: + key: str + title: str + icon: str + color: str # HEX для интерфейса + description: str + danger: bool + severity: int # чем больше, тем опаснее; влияет на выбор класса при равенстве + + @property + def bgr(self) -> tuple[int, int, int]: + h = self.color.lstrip("#") + r, g, b = (int(h[i : i + 2], 16) for i in (0, 2, 4)) + return b, g, r + + +HAZARDS: dict[str, Hazard] = { + h.key: h + for h in ( + Hazard("fire", "Пожар", "🔥", "#ef4444", "Открытое пламя · критическая опасность", True, 3), + Hazard("trash", "Сжигание мусора", "🗑️", "#eab308", "Горение бытовых отходов, пластика", True, 2), + Hazard("smoke", "Дым", "💨", "#94a3b8", "Обнаружен дым · требуется проверка", True, 1), + Hazard("factory", "Производственный дым", "🏭", "#3b82f6", "Промышленный выброс", False, 0), + Hazard("normal", "Норма", "✅", "#22c55e", "Угроз не обнаружено", False, 0), + ) +} + +_ALIASES = {"flame": "fire", "fires": "fire", "garbage": "trash", "background": "normal", "none": "normal"} + + +def normalize(name: str) -> str: + key = str(name).strip().lower() + return _ALIASES.get(key, key) + + +def get_hazard(name: str) -> Hazard: + key = normalize(name) + if key in HAZARDS: + return HAZARDS[key] + return Hazard(key, key.capitalize(), "❔", "#8b8fa3", "Неизвестный класс модели", False, 0) + + +def is_danger(name: str) -> bool: + return get_hazard(name).danger diff --git a/ecowatch/monitor.py b/ecowatch/monitor.py new file mode 100644 index 0000000..e844acd --- /dev/null +++ b/ecowatch/monitor.py @@ -0,0 +1,249 @@ +"""Фоновый мониторинг: источник → модель → трекер → журнал/уведомления/запись. + +Работает в отдельном потоке, поэтому интерфейс Streamlit остаётся отзывчивым, +а мониторинг не прерывается при переключении вкладок или перезагрузке страницы. +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections import deque +from dataclasses import dataclass, field, replace +from datetime import datetime +from pathlib import Path + +import cv2 +import numpy as np + +from .config import DetectionParams, Settings +from .detector import Detector +from .notifier import AsyncSender, TelegramNotifier, format_alert, format_incident_closed +from .overlay import draw_overlay +from .recorder import ClipRecorder +from .sources import SourceSpec +from .storage import EventStore +from .tracker import Decision, ThreatTracker + +log = logging.getLogger(__name__) + + +PREVIEW_MAX_WIDTH = 960 + + +def encode_preview(frame: np.ndarray, max_width: int = PREVIEW_MAX_WIDTH) -> bytes | None: + h, w = frame.shape[:2] + if w > max_width: + frame = cv2.resize(frame, (max_width, int(h * max_width / w)), interpolation=cv2.INTER_AREA) + ok, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) + return buf.tobytes() if ok else None + + +@dataclass +class MonitorOptions: + source: SourceSpec + params: DetectionParams = field(default_factory=DetectionParams) + notify: bool = True + notify_photo: bool = True + notify_on_clear: bool = False + notify_cooldown: float = 60.0 # мин. секунд между тревогами в Telegram + record: bool = False + record_pre: float = 3.0 + record_post: float = 10.0 + lat: float | None = None + lon: float | None = None + + +@dataclass +class MonitorState: + running: bool = False + finished: bool = False + error: str | None = None + source: str = "" + started_at: datetime | None = None + frame: np.ndarray | None = None + preview: bytes | None = None # JPEG для интерфейса, кодируется один раз в рабочем потоке + decision: Decision | None = None + fps: float = 0.0 + frames: int = 0 + incidents: int = 0 + incident_id: int | None = None + recording: bool = False + + +class Monitor: + def __init__( + self, + detector: Detector, + store: EventStore, + settings: Settings, + options: MonitorOptions, + notifier: TelegramNotifier | None = None, + sender: AsyncSender | None = None, + ) -> None: + self.detector = detector + self.store = store + self.settings = settings + self.options = options + self.notifier = notifier + self.sender = sender + self.tracker = ThreatTracker(options.params) + self._state = MonitorState(source=options.source.label) + self._lock = threading.Lock() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name="ecowatch-monitor", daemon=True) + self._incident: dict | None = None + self._last_alert = float("-inf") + + # ---------- управление ---------- + def start(self) -> None: + with self._lock: + self._state.running = True + self._state.started_at = datetime.now() + self._thread.start() + + def stop(self, wait: bool = True) -> None: + self._stop.set() + if wait and self._thread.is_alive() and threading.current_thread() is not self._thread: + self._thread.join(timeout=10) + + def update_params(self, params: DetectionParams) -> None: + """Пороги можно менять на лету, не перезапуская мониторинг.""" + self.tracker.params = params + + def state(self) -> MonitorState: + with self._lock: + return replace(self._state) + + @property + def alive(self) -> bool: + return self._thread.is_alive() + + def _set(self, **kw: object) -> None: + with self._lock: + for k, v in kw.items(): + setattr(self._state, k, v) + + # ---------- уведомления ---------- + def _notify(self, text: str, image: np.ndarray | None) -> None: + if not (self.options.notify and self.notifier and self.notifier.enabled and self.sender): + return + if image is not None and self.options.notify_photo: + self.sender.submit(self.notifier.send_photo, image, text) + else: + self.sender.submit(self.notifier.send_message, text) + + # ---------- инциденты ---------- + def _on_onset(self, d: Decision, frame: np.ndarray) -> None: + now = datetime.now() + snap: Path | None = self.settings.snapshots_dir / f"{now:%Y%m%d_%H%M%S}_{d.label}.jpg" + try: + snap.parent.mkdir(parents=True, exist_ok=True) + if not cv2.imwrite(str(snap), frame): + snap = None + except (OSError, cv2.error): + snap = None + incident_id = self.store.open_incident( + d.label, + d.confidence, + self.options.source.label, + started_at=now, + snapshot=str(snap) if snap else None, + lat=self.options.lat, + lon=self.options.lon, + ) + self._incident = {"id": incident_id, "label": d.label, "peak": d.confidence, "started": now} + log.warning("Инцидент #%s: %s (%.0f%%)", incident_id, d.label, d.confidence * 100) + if time.monotonic() - self._last_alert >= self.options.notify_cooldown: + self._last_alert = time.monotonic() + self._notify(format_alert(d.label, d.confidence, self.options.source.label, now), frame) + + def _close_incident(self) -> None: + inc = self._incident + if inc is None: + return + ended = datetime.now() + self.store.update_incident(inc["id"], ended_at=ended, peak_conf=inc["peak"], label=inc["label"]) + log.info("Инцидент #%s завершён", inc["id"]) + if self.options.notify_on_clear: + self._notify(format_incident_closed(inc["label"], inc["peak"], inc["started"], ended), None) + self._incident = None + + def _attach_clip(self, path: Path | None, incident_id: int | None) -> None: + if path and incident_id is not None: + self.store.update_incident(incident_id, clip=str(path)) + + # ---------- основной цикл ---------- + def _run(self) -> None: + source = self.options.source.build() + recorder = ( + ClipRecorder( + self.settings.clips_dir, pre_seconds=self.options.record_pre, post_seconds=self.options.record_post + ) + if self.options.record + else None + ) + clip_incident: int | None = None + fps = 0.0 + stamps: deque[float] = deque(maxlen=30) # FPS по скользящему окну последних кадров + try: + source.open() + if recorder and not source.live: + recorder.set_fps(source.fps) + while not self._stop.is_set(): + frame = source.read(timeout=5.0) + if frame is None: + if source.live: + self._set(error=getattr(source, "error", None) or "Нет кадров от источника…") + continue + self._set(finished=True) + break + + result = self.detector.predict(frame, self.tracker.params.box_conf) + d = self.tracker.update(result) + + stamps.append(time.monotonic()) + if len(stamps) > 1: + fps = (len(stamps) - 1) / max(stamps[-1] - stamps[0], 1e-6) + + annotated = draw_overlay( + frame, + d, + result.boxes, + self.options.source.label, + fps, + recording=bool(recorder and recorder.recording), + ) + + if d.onset: + self._on_onset(d, annotated) + if self._incident is not None and d.active: + self._incident["peak"] = max(self._incident["peak"], d.confidence) + self._incident["label"] = d.label + if recorder: + if source.live: + recorder.set_fps(fps) + if d.active and not recorder.recording and self._incident: + clip_incident = self._incident["id"] + self._attach_clip(recorder.push(annotated, d.active), clip_incident) + if d.cleared: + self._close_incident() + + preview = encode_preview(annotated) + with self._lock: + s = self._state + s.frame, s.preview, s.decision, s.fps, s.error = annotated, preview, d, fps, None + s.frames += 1 + s.incidents += int(d.onset) + s.incident_id = self._incident["id"] if self._incident else None + s.recording = bool(recorder and recorder.recording) + except Exception as exc: # noqa: BLE001 — ошибку показываем в интерфейсе + log.exception("Мониторинг остановлен из-за ошибки") + self._set(error=str(exc)) + finally: + self._close_incident() + if recorder: + self._attach_clip(recorder.close(), clip_incident) + source.close() + self._set(running=False, recording=False) diff --git a/ecowatch/notifier.py b/ecowatch/notifier.py new file mode 100644 index 0000000..13914ca --- /dev/null +++ b/ecowatch/notifier.py @@ -0,0 +1,140 @@ +"""Уведомления в Telegram. Отправка идёт в фоновом потоке и не тормозит видео.""" + +from __future__ import annotations + +import logging +import queue +import threading +from collections.abc import Callable +from datetime import datetime +from html import escape +from typing import Any + +import cv2 +import numpy as np +import requests + +from .hazards import get_hazard + +log = logging.getLogger(__name__) + +API = "https://api.telegram.org/bot{token}/{method}" +CAPTION_LIMIT = 1024 + + +class TelegramNotifier: + def __init__(self, token: str | None, chat_id: str | None, session: Any = None, timeout: float = 15) -> None: + self.token = token + self.chat_id = chat_id + self.session = session or requests.Session() + self.timeout = timeout + + @property + def enabled(self) -> bool: + return bool(self.token and self.chat_id) + + def _call(self, method: str, **kwargs: Any) -> bool: + if not self.enabled: + return False + try: + resp = self.session.post(API.format(token=self.token, method=method), timeout=self.timeout, **kwargs) + ok = bool(resp.json().get("ok")) + if not ok: + log.warning("Telegram %s: %s", method, resp.text[:300]) + return ok + except (requests.RequestException, ValueError) as exc: + log.warning("Telegram %s не отправлен: %s", method, exc) + return False + + def send_message(self, text: str) -> bool: + return self._call("sendMessage", json={"chat_id": self.chat_id, "text": text, "parse_mode": "HTML"}) + + def send_photo(self, image_bgr: np.ndarray, caption: str = "") -> bool: + ok, buf = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85]) + if not ok: + return self.send_message(caption) + if len(caption) > CAPTION_LIMIT: + # Подпись к фото ограничена — длинный текст отправляем отдельным сообщением. + return self.send_message(caption) and self.send_photo(image_bgr) + return self._call( + "sendPhoto", + data={"chat_id": self.chat_id, "caption": caption, "parse_mode": "HTML"}, + files={"photo": ("alert.jpg", buf.tobytes(), "image/jpeg")}, + ) + + +class AsyncSender: + """Очередь задач отправки с одним рабочим потоком.""" + + def __init__(self, maxsize: int = 32) -> None: + self._queue: queue.Queue[tuple[Callable[..., bool], tuple[Any, ...]]] = queue.Queue(maxsize) + self.sent = 0 + self.failed = 0 + self._thread = threading.Thread(target=self._run, name="telegram-sender", daemon=True) + self._thread.start() + + def submit(self, fn: Callable[..., bool], *args: Any) -> bool: + try: + self._queue.put_nowait((fn, args)) + return True + except queue.Full: + self.failed += 1 + return False + + def _run(self) -> None: + while True: + fn, args = self._queue.get() + try: + if fn(*args): + self.sent += 1 + else: + self.failed += 1 + except Exception: # noqa: BLE001 — поток отправки не должен падать + log.exception("Ошибка при отправке уведомления") + self.failed += 1 + finally: + self._queue.task_done() + + def join(self) -> None: + self._queue.join() + + +def format_alert(label: str, conf: float, source: str, when: datetime | None = None) -> str: + h = get_hazard(label) + when = when or datetime.now() + return ( + f"🚨 ЭКОЛОГИЧЕСКАЯ УГРОЗА\n\n" + f"{h.icon} {escape(h.title.upper())}\n" + f"{escape(h.description)}\n\n" + f"📊 Вероятность: {conf * 100:.0f}%\n" + f"📹 Источник: {escape(source)}\n" + f"🕐 {when:%d.%m.%Y %H:%M:%S}\n\n" + f"— EcoWatch AI" + ) + + +def format_incident_closed(label: str, peak: float, started: datetime, ended: datetime) -> str: + h = get_hazard(label) + minutes, seconds = divmod(int((ended - started).total_seconds()), 60) + return ( + f"✅ Угроза больше не наблюдается\n\n" + f"{h.icon} {escape(h.title)} · пик {peak * 100:.0f}%\n" + f"⏱ Длительность: {minutes} мин {seconds} с\n" + f"🕐 {ended:%d.%m.%Y %H:%M:%S}" + ) + + +def format_daily_report(summary: dict[str, Any]) -> str: + by_label: dict[str, int] = summary.get("by_label", {}) + lines = [f"📊 ОТЧЁТ ЗА {summary['day']:%d.%m.%Y}", ""] + if by_label: + for label, n in sorted(by_label.items(), key=lambda kv: -kv[1]): + h = get_hazard(label) + lines.append(f"{h.icon} {escape(h.title)}: {n}") + lines += ["", f"📋 Всего инцидентов: {summary['total']}"] + lines.append(f"📈 Максимальная вероятность: {summary['peak_conf'] * 100:.0f}%") + else: + lines.append("Инцидентов не зафиксировано.") + dangerous = sum(n for label, n in by_label.items() if get_hazard(label).danger) + lines += ["", "⚠️ Повышенная активность!" if dangerous > 5 else "✅ Ситуация стабильна", "", "— EcoWatch AI"] + return "\n".join(lines) diff --git a/ecowatch/overlay.py b/ecowatch/overlay.py new file mode 100644 index 0000000..fa6103c --- /dev/null +++ b/ecowatch/overlay.py @@ -0,0 +1,107 @@ +"""Отрисовка рамок и статуса поверх кадра (с поддержкой кириллицы).""" + +from __future__ import annotations + +from datetime import datetime +from functools import lru_cache + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from .detector import Box +from .hazards import get_hazard +from .tracker import Decision + +_FONT_CANDIDATES = ( + "DejaVuSans.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + "arial.ttf", + "C:/Windows/Fonts/segoeui.ttf", + "C:/Windows/Fonts/arial.ttf", + "/System/Library/Fonts/Supplemental/Arial.ttf", +) + + +@lru_cache(maxsize=16) +def _font(size: int) -> ImageFont.ImageFont | ImageFont.FreeTypeFont: + for name in _FONT_CANDIDATES: + try: + return ImageFont.truetype(name, size) + except OSError: + continue + try: + return ImageFont.load_default(size=size) + except TypeError: # Pillow < 10.1 + return ImageFont.load_default() + + +def _corners(img: np.ndarray, b: Box, color: tuple[int, int, int]) -> None: + length = max(6, min(24, (b.x2 - b.x1) // 4, (b.y2 - b.y1) // 4)) + for (x, y), (dx, dy) in ( + ((b.x1, b.y1), (1, 1)), + ((b.x2, b.y1), (-1, 1)), + ((b.x1, b.y2), (1, -1)), + ((b.x2, b.y2), (-1, -1)), + ): + cv2.line(img, (x, y), (x + dx * length, y), color, 3, cv2.LINE_AA) + cv2.line(img, (x, y), (x, y + dy * length), color, 3, cv2.LINE_AA) + + +def draw_overlay( + frame: np.ndarray, + decision: Decision, + boxes: list[Box] | None = None, + source: str = "", + fps: float | None = None, + recording: bool = False, +) -> np.ndarray: + img = frame.copy() + h, w = img.shape[:2] + scale = max(0.6, min(2.0, w / 960)) + fs = int(15 * scale) + bar_h = int(34 * scale) + hz = get_hazard(decision.label) + texts: list[tuple[tuple[int, int], str, tuple[int, int, int], int]] = [] + + for b in boxes or []: + bh = get_hazard(b.label) + cv2.rectangle(img, (b.x1, b.y1), (b.x2, b.y2), bh.bgr, 1, cv2.LINE_AA) + _corners(img, b, bh.bgr) + text = f"{bh.title} {b.conf * 100:.0f}%" + tw = int(_font(fs).getlength(text)) + ty = max(0, b.y1 - fs - 8) + cv2.rectangle(img, (b.x1, ty), (b.x1 + tw + 10, ty + fs + 8), bh.bgr, -1) + texts.append(((b.x1 + 5, ty + 3), text, (255, 255, 255), fs)) + + # Полупрозрачная верхняя панель статуса + top = img[:bar_h].astype(np.float32) + img[:bar_h] = (top * 0.35).astype(np.uint8) + dot = (60, 60, 235) if decision.active else ((40, 190, 235) if decision.frame_danger else (80, 200, 80)) + cv2.circle(img, (int(16 * scale), bar_h // 2), int(6 * scale), dot, -1, cv2.LINE_AA) + + status = f"{hz.title.upper()} · {decision.confidence * 100:.0f}%" + if not decision.active and decision.frame_danger: + status = "ПРОВЕРКА…" + texts.append(((int(30 * scale), bar_h // 2 - fs // 2 - 2), status, (235, 235, 235), fs)) + right = datetime.now().strftime("%H:%M:%S") + if fps is not None: + right = f"{fps:.1f} FPS · {right}" + if source: + right = f"{source} · {right}" + rw = int(_font(fs - 2).getlength(right)) + texts.append(((w - rw - 12, bar_h // 2 - fs // 2), right, (170, 170, 170), fs - 2)) + + if recording: + cv2.circle(img, (w - 18, bar_h + 18), 7, (40, 40, 230), -1, cv2.LINE_AA) + texts.append(((w - 58, bar_h + 9), "REC", (40, 40, 230), fs - 3)) + + if decision.active and not boxes: + cv2.rectangle(img, (2, bar_h + 2), (w - 3, h - 3), hz.bgr, max(2, int(3 * scale))) + + # Весь текст — за одно преобразование в PIL (раньше было по одному на каждую надпись). + pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) + draw = ImageDraw.Draw(pil) + for pos, text, (b, g, r), size in texts: + draw.text(pos, text, font=_font(size), fill=(r, g, b)) + return cv2.cvtColor(np.asarray(pil), cv2.COLOR_RGB2BGR) diff --git a/ecowatch/recorder.py b/ecowatch/recorder.py new file mode 100644 index 0000000..34a1bda --- /dev/null +++ b/ecowatch/recorder.py @@ -0,0 +1,87 @@ +"""Запись видеоклипов инцидентов с предзаписью (несколько секунд до тревоги).""" + +from __future__ import annotations + +import time +from collections import deque +from datetime import datetime +from pathlib import Path + +import cv2 +import numpy as np + + +class ClipRecorder: + def __init__( + self, + out_dir: str | Path, + fps: float = 10.0, + pre_seconds: float = 3.0, + post_seconds: float = 10.0, + max_seconds: float = 300.0, + ) -> None: + self.out_dir = Path(out_dir) + self.fps = max(1.0, float(fps)) + self.pre_seconds = pre_seconds + self.post_seconds = post_seconds + self.max_seconds = max_seconds + # Буфер с запасом на 30 FPS; при старте записи берём из него последние pre_seconds. + self._buffer: deque[np.ndarray] = deque(maxlen=max(1, int(30 * pre_seconds))) + self._writer: cv2.VideoWriter | None = None + self._size: tuple[int, int] | None = None + self._path: Path | None = None + self._started = 0.0 + self._last_active = 0.0 + + @property + def recording(self) -> bool: + return self._writer is not None + + def set_fps(self, fps: float) -> None: + """FPS клипа = реальная скорость обработки; меняется только между записями.""" + if self._writer is None and fps > 0: + self.fps = max(1.0, min(30.0, float(fps))) + + def _open(self, frame: np.ndarray) -> None: + self.out_dir.mkdir(parents=True, exist_ok=True) + h, w = frame.shape[:2] + self._size = (w, h) + self._path = self.out_dir / f"incident_{datetime.now():%Y%m%d_%H%M%S}.mp4" + self._writer = cv2.VideoWriter(str(self._path), cv2.VideoWriter_fourcc(*"mp4v"), self.fps, self._size) + if not self._writer.isOpened(): + self._writer = None + raise RuntimeError(f"Не удалось открыть файл записи {self._path}") + + def _write(self, frame: np.ndarray) -> None: + assert self._writer is not None and self._size is not None + if (frame.shape[1], frame.shape[0]) != self._size: + frame = cv2.resize(frame, self._size) + self._writer.write(frame) + + def push(self, frame: np.ndarray, active: bool, now: float | None = None) -> Path | None: + """Добавляет кадр. Возвращает путь к клипу, когда тот закончен.""" + now = time.monotonic() if now is None else now + if active: + self._last_active = now + if self._writer is None: + self._open(frame) + self._started = now + keep = int(self.fps * self.pre_seconds) + for old in list(self._buffer)[-keep:] if keep > 0 else []: + self._write(old) + self._buffer.clear() + if self._writer is None: + self._buffer.append(frame) + return None + self._write(frame) + if now - self._last_active > self.post_seconds or now - self._started > self.max_seconds: + return self.close() + return None + + def close(self) -> Path | None: + if self._writer is None: + return None + self._writer.release() + self._writer = None + path, self._path = self._path, None + return path diff --git a/ecowatch/sources.py b/ecowatch/sources.py new file mode 100644 index 0000000..e999054 --- /dev/null +++ b/ecowatch/sources.py @@ -0,0 +1,211 @@ +"""Источники кадров: веб-камера, IP-камера (RTSP/HTTP), видеофайл, экран.""" + +from __future__ import annotations + +import threading +import time +from abc import ABC, abstractmethod +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + + +class FrameSource(ABC): + live: bool = True + fps: float = 0.0 + + @abstractmethod + def open(self) -> None: ... + + @abstractmethod + def read(self, timeout: float = 5.0) -> np.ndarray | None: + """Возвращает следующий кадр (BGR) или None, если кадров больше нет/нет связи.""" + + @abstractmethod + def close(self) -> None: ... + + def __enter__(self) -> FrameSource: + self.open() + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + +class VideoFileSource(FrameSource): + """Видеофайл читается последовательно — каждый кадр будет проанализирован.""" + + live = False + + def __init__(self, path: str, frame_step: int = 1) -> None: + self.path = path + self.frame_step = max(1, frame_step) + self._cap: cv2.VideoCapture | None = None + + def open(self) -> None: + self._cap = cv2.VideoCapture(self.path) + if not self._cap.isOpened(): + raise RuntimeError(f"Не удалось открыть видео: {self.path}") + self.fps = (self._cap.get(cv2.CAP_PROP_FPS) or 25.0) / self.frame_step + + def read(self, timeout: float = 5.0) -> np.ndarray | None: + assert self._cap is not None + for _ in range(self.frame_step - 1): + if not self._cap.grab(): + return None + ok, frame = self._cap.read() + return frame if ok else None + + def close(self) -> None: + if self._cap is not None: + self._cap.release() + self._cap = None + + +class StreamSource(FrameSource): + """Камера/поток. Кадры читаются фоновым потоком, отдаётся всегда самый свежий — + поэтому медленная модель не создаёт растущую задержку.""" + + live = True + + def __init__(self, target: int | str, reconnect: bool = True) -> None: + self.target = target + self.reconnect = reconnect + self._cap: cv2.VideoCapture | None = None + self._frame: np.ndarray | None = None + self._seq = 0 + self._read_seq = 0 + self._cond = threading.Condition() + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self.error: str | None = None + + def _connect(self) -> bool: + cap = cv2.VideoCapture(self.target) + if not cap.isOpened(): + cap.release() + return False + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + self._cap = cap + self.fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 + return True + + def open(self) -> None: + if not self._connect(): + raise RuntimeError(f"Не удалось подключиться к источнику: {self.target}") + self._stop.clear() + self._thread = threading.Thread(target=self._grab_loop, name="frame-grabber", daemon=True) + self._thread.start() + + def _grab_loop(self) -> None: + backoff = 1.0 + while not self._stop.is_set(): + cap = self._cap + ok, frame = cap.read() if cap is not None else (False, None) + if ok: + backoff = 1.0 + with self._cond: + self._frame = frame + self._seq += 1 + self._cond.notify_all() + continue + if not self.reconnect: + break + self.error = "Потеряна связь с камерой, переподключение…" + if cap is not None: + cap.release() + self._cap = None + if self._stop.wait(backoff): + break + backoff = min(backoff * 2, 30.0) + if self._connect(): + self.error = None + with self._cond: + self._cond.notify_all() + + def read(self, timeout: float = 5.0) -> np.ndarray | None: + deadline = time.monotonic() + timeout + with self._cond: + while self._seq == self._read_seq and not self._stop.is_set(): + remaining = deadline - time.monotonic() + if remaining <= 0 or (self._thread is not None and not self._thread.is_alive()): + return None + self._cond.wait(remaining) + if self._seq == self._read_seq: + return None + self._read_seq = self._seq + return self._frame + + def close(self) -> None: + self._stop.set() + with self._cond: + self._cond.notify_all() + if self._thread is not None: + self._thread.join(timeout=3) + if self._cap is not None: + self._cap.release() + self._cap = None + + +class ScreenSource(FrameSource): + live = True + + def __init__(self, monitor: int = 1, max_fps: float = 10.0) -> None: + self.monitor = monitor + self.fps = max_fps + self._sct = None + self._last = 0.0 + + def open(self) -> None: + from mss import mss + + self._sct = mss() + if self.monitor >= len(self._sct.monitors): + raise RuntimeError(f"Монитор {self.monitor} не найден") + + def read(self, timeout: float = 5.0) -> np.ndarray | None: + assert self._sct is not None + wait = 1.0 / self.fps - (time.monotonic() - self._last) + if wait > 0: + time.sleep(wait) + self._last = time.monotonic() + img = np.asarray(self._sct.grab(self._sct.monitors[self.monitor])) + return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) + + def close(self) -> None: + if self._sct is not None: + self._sct.close() + self._sct = None + + +@dataclass +class SourceSpec: + kind: str # "webcam" | "stream" | "file" | "screen" + value: str | int + label: str + + def build(self) -> FrameSource: + if self.kind == "webcam": + return StreamSource(int(self.value)) + if self.kind == "stream": + return StreamSource(str(self.value)) + if self.kind == "file": + return VideoFileSource(str(self.value)) + if self.kind == "screen": + return ScreenSource(int(self.value)) + raise ValueError(f"Неизвестный тип источника: {self.kind}") + + +def parse_source(value: str) -> SourceSpec: + """Для CLI: '0' → веб-камера, 'rtsp://…' → поток, 'screen' / 'screen:2' → экран, иначе — файл.""" + v = value.strip() + if v.isdigit(): + return SourceSpec("webcam", int(v), f"Камера {v}") + if v.startswith("screen"): + idx = int(v.split(":", 1)[1]) if ":" in v else 1 + return SourceSpec("screen", idx, f"Экран {idx}") + if "://" in v: + return SourceSpec("stream", v, v.split("@")[-1]) + return SourceSpec("file", v, Path(v).name) diff --git a/ecowatch/storage.py b/ecowatch/storage.py new file mode 100644 index 0000000..efc4830 --- /dev/null +++ b/ecowatch/storage.py @@ -0,0 +1,136 @@ +"""Журнал инцидентов в SQLite (вместо перезаписываемого JSON-файла).""" + +from __future__ import annotations + +import json +import sqlite3 +import threading +from collections import Counter +from datetime import date, datetime +from pathlib import Path +from typing import Any + +from .hazards import normalize + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS incidents ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + started_at TEXT NOT NULL, + ended_at TEXT, + label TEXT NOT NULL, + peak_conf REAL NOT NULL, + source TEXT NOT NULL DEFAULT '', + snapshot TEXT, + clip TEXT, + lat REAL, + lon REAL +); +CREATE INDEX IF NOT EXISTS idx_incidents_started ON incidents(started_at); +""" + +COLUMNS = ("id", "started_at", "ended_at", "label", "peak_conf", "source", "snapshot", "clip", "lat", "lon") + + +def _ts(dt: datetime | None) -> str | None: + return dt.isoformat(timespec="seconds") if dt else None + + +class EventStore: + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._lock = threading.Lock() + self._conn = sqlite3.connect(self.path, check_same_thread=False) + self._conn.row_factory = sqlite3.Row + with self._lock, self._conn: + self._conn.executescript(_SCHEMA) + + def close(self) -> None: + with self._lock: + self._conn.close() + + def open_incident( + self, + label: str, + conf: float, + source: str, + started_at: datetime | None = None, + snapshot: str | None = None, + lat: float | None = None, + lon: float | None = None, + ) -> int: + with self._lock, self._conn: + cur = self._conn.execute( + "INSERT INTO incidents (started_at, label, peak_conf, source, snapshot, lat, lon)" + " VALUES (?, ?, ?, ?, ?, ?, ?)", + (_ts(started_at or datetime.now()), label, float(conf), source, snapshot, lat, lon), + ) + return int(cur.lastrowid) + + def update_incident(self, incident_id: int, **fields: Any) -> None: + allowed = {"ended_at", "label", "peak_conf", "snapshot", "clip"} + unknown = set(fields) - allowed + if unknown: + raise ValueError(f"Нельзя обновить поля: {unknown}") + if not fields: + return + if isinstance(fields.get("ended_at"), datetime): + fields["ended_at"] = _ts(fields["ended_at"]) + assignments = ", ".join(f"{k} = ?" for k in fields) + with self._lock, self._conn: + self._conn.execute(f"UPDATE incidents SET {assignments} WHERE id = ?", (*fields.values(), incident_id)) + + def list(self, since: datetime | None = None, limit: int | None = None) -> list[dict[str, Any]]: + sql = "SELECT * FROM incidents" + args: list[Any] = [] + if since: + sql += " WHERE started_at >= ?" + args.append(_ts(since)) + sql += " ORDER BY started_at DESC, id DESC" + if limit: + sql += " LIMIT ?" + args.append(limit) + with self._lock: + return [dict(r) for r in self._conn.execute(sql, args).fetchall()] + + def count(self) -> int: + with self._lock: + return int(self._conn.execute("SELECT COUNT(*) FROM incidents").fetchone()[0]) + + def day_summary(self, day: date | None = None) -> dict[str, Any]: + day = day or date.today() + prefix = day.isoformat() + with self._lock: + rows = self._conn.execute( + "SELECT label, peak_conf FROM incidents WHERE substr(started_at, 1, 10) = ?", (prefix,) + ).fetchall() + counts = Counter(r["label"] for r in rows) + peak = max((r["peak_conf"] for r in rows), default=0.0) + return {"day": day, "total": len(rows), "by_label": dict(counts), "peak_conf": peak} + + def clear(self) -> None: + with self._lock, self._conn: + self._conn.execute("DELETE FROM incidents") + + def import_legacy_json(self, path: str | Path) -> int: + """Переносит detection_history.json из старой версии. Возвращает число записей.""" + path = Path(path) + if not path.is_file(): + return 0 + try: + entries = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return 0 + rows = [] + for e in entries if isinstance(entries, list) else []: + try: + started = datetime.strptime(str(e["time"]), "%Y-%m-%d %H:%M:%S") + except (KeyError, ValueError): + continue + conf = float(e.get("confidence", 0)) / 100.0 # в старом формате — проценты + rows.append((_ts(started), normalize(e.get("label", "normal")), conf, str(e.get("source", "")))) + with self._lock, self._conn: + self._conn.executemany( + "INSERT INTO incidents (started_at, label, peak_conf, source) VALUES (?, ?, ?, ?)", rows + ) + return len(rows) diff --git a/ecowatch/tracker.py b/ecowatch/tracker.py new file mode 100644 index 0000000..d47c107 --- /dev/null +++ b/ecowatch/tracker.py @@ -0,0 +1,104 @@ +"""Временное сглаживание и гистерезис: превращает поток кадров в инциденты. + +Вероятности сглаживаются экспоненциальным средним, а инцидент начинается +только после `confirm_frames` подряд опасных кадров и завершается после +`clear_frames` подряд чистых. Это убирает «мигание» тревоги и ложные +срабатывания на единичных кадрах. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from .config import DetectionParams +from .detector import FrameResult +from .hazards import get_hazard + + +@dataclass +class Decision: + label: str # активный класс угрозы или "normal" + confidence: float # сглаженная вероятность label + active: bool # идёт ли сейчас инцидент + frame_danger: bool # опасен ли текущий кадр (до гистерезиса) + scores: dict[str, float] = field(default_factory=dict) + onset: bool = False # инцидент начался на этом кадре + cleared: bool = False # инцидент завершился на этом кадре + + @property + def ambiguous(self) -> bool: + ranked = sorted(self.scores.values(), reverse=True) + return len(ranked) >= 2 and ranked[0] - ranked[1] < 0.12 + + +def classify_frame(scores: dict[str, float], params: DetectionParams) -> tuple[str, float, bool]: + """Решение по одному набору вероятностей: (кандидат, вероятность, опасно ли).""" + if not scores: + return "normal", 0.0, False + ranked = sorted(scores.items(), key=lambda kv: (kv[1], get_hazard(kv[0]).severity), reverse=True) + top_label, top_conf = ranked[0] + second = ranked[1][1] if len(ranked) > 1 else 0.0 + danger = get_hazard(top_label).danger and top_conf >= params.conf_threshold and top_conf - second >= params.margin + return top_label, top_conf, danger + + +class ThreatTracker: + def __init__(self, params: DetectionParams | None = None) -> None: + self.params = params or DetectionParams() + self.reset() + + def reset(self) -> None: + self._ema: dict[str, float] = {} + self._danger_streak = 0 + self._clear_streak = 0 + self.active = False + self.label = "normal" + + def _smooth(self, scores: dict[str, float]) -> dict[str, float]: + a = self.params.ema_alpha + if not self._ema: + self._ema = dict(scores) + else: + for k in set(self._ema) | set(scores): + self._ema[k] = a * scores.get(k, 0.0) + (1 - a) * self._ema.get(k, 0.0) + return dict(self._ema) + + def update(self, result: FrameResult) -> Decision: + p = self.params + smoothed = self._smooth(result.scores) + candidate, cand_conf, frame_danger = classify_frame(smoothed, p) + onset = cleared = False + + if frame_danger: + self._danger_streak += 1 + self._clear_streak = 0 + else: + self._clear_streak += 1 + self._danger_streak = 0 + + if not self.active: + if self._danger_streak >= p.confirm_frames: + self.active, onset, self.label = True, True, candidate + else: + if frame_danger and get_hazard(candidate).severity >= get_hazard(self.label).severity: + self.label = candidate # эскалация: дым → пожар + if self._clear_streak >= p.clear_frames: + self.active, cleared = False, True + + if self.active: + label, conf = self.label, smoothed.get(self.label, 0.0) + else: + # Пока угроза не подтверждена, статус — «норма» (frame_danger покажет, что идёт проверка). + label = "normal" + conf = smoothed.get("normal", 1.0 - cand_conf if get_hazard(candidate).danger else cand_conf) + return Decision(label, conf, self.active, frame_danger, smoothed, onset, cleared) + + +def evaluate_single(result: FrameResult, params: DetectionParams | None = None) -> Decision: + """Решение для отдельного изображения — без сглаживания и гистерезиса.""" + params = params or DetectionParams() + candidate, conf, danger = classify_frame(result.scores, params) + label = candidate if danger else "normal" + if not danger: + conf = result.scores.get("normal", conf) + return Decision(label, conf, danger, danger, dict(result.scores), onset=danger) diff --git a/ecowatch/ui.py b/ecowatch/ui.py new file mode 100644 index 0000000..aeaaa73 --- /dev/null +++ b/ecowatch/ui.py @@ -0,0 +1,495 @@ +"""Веб-интерфейс EcoWatch AI на Streamlit.""" + +from __future__ import annotations + +import threading +from datetime import datetime, timedelta +from pathlib import Path + +import cv2 +import numpy as np +import pandas as pd +import plotly.express as px +import plotly.graph_objects as go +import streamlit as st + +from .config import ROOT_DIR, DetectionParams, Settings +from .detector import Detector +from .hazards import HAZARDS, get_hazard +from .monitor import Monitor, MonitorOptions +from .notifier import AsyncSender, TelegramNotifier, format_daily_report +from .overlay import draw_overlay +from .sources import SourceSpec +from .storage import EventStore +from .tracker import Decision, evaluate_single + +CSS = """ + +""" + +SOURCES = {"webcam": "Веб-камера", "stream": "IP-камера (RTSP/HTTP)", "file": "Видеофайл", "screen": "Захват экрана"} + + +# ---------- общие ресурсы (одни на весь сервер) ---------- +@st.cache_resource +def get_settings() -> Settings: + s = Settings.from_env() + s.ensure_dirs() + return s + + +@st.cache_resource(show_spinner="Загрузка модели…") +def get_detector(path: str, device: str | None) -> Detector: + return Detector(path, device) + + +@st.cache_resource +def get_store(path: str) -> EventStore: + store = EventStore(path) + legacy = ROOT_DIR / "detection_history.json" + if store.count() == 0 and legacy.is_file(): + store.import_legacy_json(legacy) + return store + + +@st.cache_resource +def get_sender() -> AsyncSender: + return AsyncSender() + + +class MonitorHolder: + """Один активный мониторинг на сервер — переживает перезагрузку страницы.""" + + def __init__(self) -> None: + self.monitor: Monitor | None = None + self.lock = threading.Lock() + + +@st.cache_resource +def get_holder() -> MonitorHolder: + return MonitorHolder() + + +# ---------- вспомогательное ---------- +def scores_html(scores: dict[str, float]) -> str: + rows = [] + for name, p in sorted(scores.items(), key=lambda kv: -kv[1]): + h = get_hazard(name) + rows.append( + f"
" + f"{h.icon} {h.title}{p * 100:.1f}%
" + f"
" + ) + return "".join(rows) + + +def status_card(d: Decision | None, subtitle: str = "") -> str: + if d is None: + return ( + "
Система готова
" + "
Выберите источник и запустите мониторинг
" + ) + h = get_hazard(d.label) + if d.active: + cls, title, desc = "ew-danger", f"{h.icon} {h.title.upper()}", h.description + elif d.frame_danger: + cls, title, desc = "ew-warn", "⏳ Проверка…", "Возможная угроза, ждём подтверждения на нескольких кадрах" + else: + cls, title, desc = "ew-safe", "✅ Угроз не обнаружено", "" + note = "
Результат неоднозначен
" if d.ambiguous else "" + return ( + f"
{title}
" + f"
{desc}
{scores_html(d.scores)}
" + f"{note}
{subtitle}
" + ) + + +def incidents_df(store: EventStore) -> pd.DataFrame: + df = pd.DataFrame(store.list()) + if df.empty: + return df + df["started_at"] = pd.to_datetime(df["started_at"]) + df["ended_at"] = pd.to_datetime(df["ended_at"]) + df["Тип"] = df["label"].map(lambda k: f"{get_hazard(k).icon} {get_hazard(k).title}") + df["Вероятность, %"] = (df["peak_conf"] * 100).round(1) + df["Длительность, с"] = (df["ended_at"] - df["started_at"]).dt.total_seconds() + return df + + +def existing_file(value: object) -> Path | None: + """Путь из журнала, если файл ещё на диске (пустые значения pandas приходят как NaN).""" + if isinstance(value, str) and value and Path(value).is_file(): + return Path(value) + return None + + +def empty_state(icon: str, text: str) -> None: + st.markdown( + f"
" + f"
{icon}
{text}
", + unsafe_allow_html=True, + ) + + +# ---------- сайдбар ---------- +def sidebar( + settings: Settings, detector: Detector, holder: MonitorHolder, store: EventStore +) -> tuple[MonitorOptions | None, DetectionParams]: + """Рисует настройки. Возвращает (опции запуска, если нажали «Запустить»; текущие пороги).""" + sb = st.sidebar + logo = ROOT_DIR / "logo" / "logo_small.png" + if logo.is_file(): + sb.image(str(logo), width=130) + sb.markdown("### 🛡️ EcoWatch AI") + sb.caption(f"Модель: `{detector.path.name}` · {detector.task} · классы: {', '.join(detector.names)}") + + running = holder.monitor is not None and holder.monitor.alive + sb.subheader("Источник") + kind = sb.selectbox("Тип", list(SOURCES), format_func=SOURCES.get, disabled=running) + spec: SourceSpec | None = None + if kind == "webcam": + cam = sb.number_input("Номер камеры", 0, 10, 0, disabled=running) + spec = SourceSpec("webcam", int(cam), f"Камера {cam}") + elif kind == "stream": + url = sb.text_input("URL потока", placeholder="rtsp://user:pass@192.168.1.10:554/stream", disabled=running) + if url.strip(): + spec = SourceSpec("stream", url.strip(), url.split("@")[-1][:40]) + elif kind == "file": + up = sb.file_uploader("Видео", type=["mp4", "avi", "mov", "mkv"], disabled=running) + if up is not None: + path = settings.data_dir / "uploads" / up.name + path.parent.mkdir(parents=True, exist_ok=True) + if not path.exists() or path.stat().st_size != up.size: + path.write_bytes(up.getbuffer()) + spec = SourceSpec("file", str(path), up.name) + else: + mon = sb.number_input("Монитор", 1, 8, 1, disabled=running) + spec = SourceSpec("screen", int(mon), f"Экран {mon}") + + with sb.expander("Чувствительность", expanded=False): + p = DetectionParams() + params = DetectionParams( + conf_threshold=st.slider( + "Порог вероятности", + 0.3, + 0.99, + p.conf_threshold, + 0.01, + help="Минимальная сглаженная вероятность опасного класса", + ), + margin=st.slider("Отрыв от 2-го класса", 0.0, 0.6, p.margin, 0.01, help="Защита от неоднозначных кадров"), + box_conf=st.slider("Порог рамки (detect)", 0.1, 0.9, p.box_conf, 0.05), + ema_alpha=st.slider("Сглаживание (вес нового кадра)", 0.05, 1.0, p.ema_alpha, 0.05), + confirm_frames=st.slider("Кадров для подтверждения", 1, 30, p.confirm_frames), + clear_frames=st.slider("Кадров для отбоя", 1, 120, p.clear_frames), + ) + if running: + holder.monitor.update_params(params) + + with sb.expander("Telegram", expanded=False): + if settings.telegram_configured: + st.success("Бот настроен") + else: + st.warning("Не настроен. Укажите ECOWATCH_TG_TOKEN и ECOWATCH_TG_CHAT_ID в файле .env") + notify = st.checkbox("Отправлять тревоги", True, disabled=not settings.telegram_configured) + notify_photo = st.checkbox("Прикладывать снимок", True) + notify_on_clear = st.checkbox("Сообщать об окончании угрозы", False) + cooldown = st.slider("Пауза между тревогами, с", 0, 600, 60, 10) + c1, c2 = st.columns(2) + tg = TelegramNotifier(settings.telegram_token, settings.telegram_chat_id) + if c1.button("Тест", disabled=not settings.telegram_configured, width="stretch"): + ok = tg.send_message(f"✅ EcoWatch AI на связи\n{datetime.now():%d.%m.%Y %H:%M:%S}") + (st.success if ok else st.error)("Отправлено" if ok else "Ошибка отправки") + if c2.button("Отчёт", disabled=not settings.telegram_configured, width="stretch"): + ok = tg.send_message(format_daily_report(store.day_summary())) + (st.success if ok else st.error)("Отправлен" if ok else "Ошибка отправки") + + with sb.expander("Запись и карта", expanded=False): + record = st.checkbox("Записывать клипы инцидентов", False) + pre = st.slider("Секунд до тревоги", 0, 10, 3, disabled=not record) + post = st.slider("Секунд после окончания", 2, 60, 10, disabled=not record) + lat = st.number_input("Широта камеры", value=settings.camera_lat or 0.0, format="%.5f") + lon = st.number_input("Долгота камеры", value=settings.camera_lon or 0.0, format="%.5f") + has_geo = bool(lat or lon) + + sb.divider() + if running: + if sb.button("⏹ Остановить", type="primary", width="stretch"): + holder.monitor.stop() + st.rerun() + return None, params + if sb.button("▶ Запустить мониторинг", type="primary", width="stretch", disabled=spec is None): + return MonitorOptions( + source=spec, # type: ignore[arg-type] + params=params, + notify=notify and settings.telegram_configured, + notify_photo=notify_photo, + notify_on_clear=notify_on_clear, + notify_cooldown=float(cooldown), + record=record, + record_pre=float(pre), + record_post=float(post), + lat=lat if has_geo else None, + lon=lon if has_geo else None, + ), params + return None, params + + +# ---------- вкладки ---------- +@st.fragment(run_every=0.25) +def live_view(holder: MonitorHolder) -> None: + mon = holder.monitor + col_v, col_s = st.columns([3, 1.2]) + if mon is None: + col_v.markdown("
📹  Мониторинг не запущен
", unsafe_allow_html=True) + col_s.markdown(status_card(None), unsafe_allow_html=True) + return + s = mon.state() + if s.preview is not None: + col_v.image(s.preview, width="stretch") + else: + col_v.markdown("
⏳ Подключение к источнику…
", unsafe_allow_html=True) + if s.error: + col_v.error(s.error) + if s.finished: + col_v.info("Видео обработано до конца.") + elif not s.running: + col_v.info("Мониторинг остановлен.") + + col_s.markdown(status_card(s.decision, f"Источник: {s.source}"), unsafe_allow_html=True) + m1, m2 = col_s.columns(2) + m1.metric("FPS", f"{s.fps:.1f}") + m2.metric("Кадров", s.frames) + m3, m4 = col_s.columns(2) + m3.metric("Инцидентов", s.incidents) + elapsed = (datetime.now() - s.started_at) if s.started_at else timedelta() + m4.metric("Время", str(elapsed).split(".")[0]) + if s.recording: + col_s.markdown("🔴 **Идёт запись клипа**") + + +def tab_image(detector: Detector, params: DetectionParams) -> None: + st.caption("Проверка отдельного снимка — без сглаживания по времени.") + up = st.file_uploader("Изображение", type=["jpg", "jpeg", "png", "webp", "bmp"], key="img_upload") + if up is None: + empty_state("🖼️", "Загрузите изображение для анализа") + return + frame = cv2.imdecode(np.frombuffer(up.getvalue(), np.uint8), cv2.IMREAD_COLOR) + if frame is None: + st.error("Не удалось прочитать изображение") + return + result = detector.predict(frame, params.box_conf) + d = evaluate_single(result, params) + c1, c2 = st.columns([3, 1.2]) + c1.image(draw_overlay(frame, d, result.boxes, up.name), channels="BGR", width="stretch") + c2.markdown(status_card(d), unsafe_allow_html=True) + + +def tab_analytics(df: pd.DataFrame) -> None: + if df.empty: + empty_state("📊", "Аналитика появится после первых инцидентов") + return + period = st.radio("Период", ["24 часа", "7 дней", "30 дней", "Всё время"], index=1, horizontal=True) + days = {"24 часа": 1, "7 дней": 7, "30 дней": 30}.get(period) + if days: + df = df[df["started_at"] >= datetime.now() - timedelta(days=days)] + if df.empty: + empty_state("📊", "За выбранный период инцидентов нет") + return + + cols = st.columns(4) + cols[0].metric("Инцидентов", len(df)) + cols[1].metric("🔥 Пожаров", int((df["label"] == "fire").sum())) + cols[2].metric("💨 Дыма", int((df["label"] == "smoke").sum())) + cols[3].metric("Ср. длительность", f"{df['Длительность, с'].mean():.0f} с" if df["ended_at"].notna().any() else "—") + + colors = {f"{h.icon} {h.title}": h.color for h in HAZARDS.values()} + layout = dict(margin=dict(t=40, b=10, l=10, r=10), legend=dict(orientation="h", y=-0.2)) + + c1, c2 = st.columns(2) + bucket = df["started_at"].dt.floor("h" if days == 1 else "D") + by_time = df.assign(t=bucket).groupby(["t", "Тип"]).size().reset_index(name="n") + by_time["t"] = by_time["t"].dt.strftime("%d.%m %H:00" if days == 1 else "%d.%m.%Y") + fig = px.bar( + by_time, + x="t", + y="n", + color="Тип", + color_discrete_map=colors, + labels={"t": "", "n": ""}, + title="Инциденты по времени", + ) + c1.plotly_chart(fig.update_layout(**layout), width="stretch") + + by_type = df["Тип"].value_counts().reset_index() + by_type.columns = ["Тип", "n"] + fig = px.pie( + by_type, values="n", names="Тип", hole=0.5, color="Тип", color_discrete_map=colors, title="Распределение" + ) + c2.plotly_chart(fig.update_layout(**layout), width="stretch") + + c3, c4 = st.columns(2) + fig = px.histogram( + df, + x="Вероятность, %", + color="Тип", + nbins=20, + color_discrete_map=colors, + title="Пиковая вероятность", + labels={"count": ""}, + ) + c3.plotly_chart(fig.update_layout(**layout, bargap=0.05), width="stretch") + + days_ru = ["Пн", "Вт", "Ср", "Чт", "Пт", "Сб", "Вс"] + heat = np.zeros((7, 24), dtype=int) + for ts in df["started_at"]: + heat[ts.weekday(), ts.hour] += 1 + fig = go.Figure( + go.Heatmap( + z=heat, + x=[f"{h}:00" for h in range(24)], + y=days_ru, + colorscale="YlOrRd", + hovertemplate="%{y} %{x}
Инцидентов: %{z}", + ) + ) + c4.plotly_chart( + fig.update_layout( + title="Активность по часам", margin=dict(t=40, b=10, l=10, r=10), yaxis=dict(autorange="reversed") + ), + width="stretch", + ) + + +def tab_map(df: pd.DataFrame) -> None: + geo = df.dropna(subset=["lat", "lon"]) if not df.empty else df + if geo.empty: + empty_state( + "🗺️", "Нет инцидентов с координатами. Укажите широту и долготу камеры в настройках «Запись и карта»." + ) + return + points = geo.groupby(["lat", "lon"]).agg(n=("id", "count"), last=("started_at", "max")).reset_index() + fig = px.scatter_map( + points, + lat="lat", + lon="lon", + size="n", + hover_data={"n": True, "last": True}, + zoom=11, + height=520, + color_discrete_sequence=["#ef4444"], + ) + st.plotly_chart(fig.update_layout(margin=dict(t=0, b=0, l=0, r=0)), width="stretch") + st.caption(f"Камер: {len(points)} · инцидентов: {int(points['n'].sum())}") + + +def tab_journal(df: pd.DataFrame, store: EventStore) -> None: + if df.empty: + empty_state("📋", "Журнал пуст") + return + c1, c2 = st.columns([2, 1]) + types = c1.multiselect("Тип", sorted(df["Тип"].unique()), default=sorted(df["Тип"].unique())) + min_conf = c2.slider("Мин. вероятность, %", 0, 100, 0) + view = df[df["Тип"].isin(types) & (df["Вероятность, %"] >= min_conf)] + table = view[["id", "started_at", "ended_at", "Тип", "Вероятность, %", "Длительность, с", "source"]].rename( + columns={"id": "№", "started_at": "Начало", "ended_at": "Конец", "source": "Источник"} + ) + event = st.dataframe( + table, width="stretch", height=360, hide_index=True, on_select="rerun", selection_mode="single-row" + ) + + rows = event.selection.rows if event and event.selection else [] + if rows: + rec = view.iloc[rows[0]] + st.markdown(f"#### Инцидент №{rec['id']} · {rec['Тип']}") + a, b = st.columns(2) + snapshot, clip = existing_file(rec["snapshot"]), existing_file(rec["clip"]) + if snapshot: + a.image(str(snapshot), caption="Снимок в момент тревоги", width="stretch") + else: + a.caption("Снимок недоступен") + if clip: + b.video(str(clip)) + b.download_button("Скачать клип", clip.read_bytes(), clip.name) + else: + b.caption("Клип не записывался") + else: + st.caption("Выберите строку, чтобы посмотреть снимок и клип.") + + c1, c2, c3 = st.columns(3) + c1.download_button( + "⬇ CSV", + table.to_csv(index=False).encode("utf-8-sig"), + f"ecowatch_{datetime.now():%Y%m%d}.csv", + "text/csv", + width="stretch", + ) + c2.download_button( + "⬇ Отчёт за сегодня", + format_daily_report(store.day_summary()) + .replace("", "") + .replace("", "") + .replace("", "") + .replace("", "") + .encode("utf-8"), + f"report_{datetime.now():%Y%m%d}.txt", + width="stretch", + ) + with c3.popover("🗑 Очистить журнал", width="stretch"): + st.write("Удалить все записи? Снимки и клипы на диске останутся.") + if st.button("Да, удалить", type="primary"): + store.clear() + st.rerun() + + +# ---------- точка входа ---------- +def main() -> None: + st.set_page_config(page_title="EcoWatch AI", page_icon="🛡️", layout="wide") + st.markdown(CSS, unsafe_allow_html=True) + + settings = get_settings() + try: + detector = get_detector(str(settings.model_path), settings.device) + except Exception as exc: # noqa: BLE001 + st.error(f"Не удалось загрузить модель: {exc}") + st.stop() + store = get_store(str(settings.db_path)) + holder = get_holder() + + options, params = sidebar(settings, detector, holder, store) + if options is not None: + with holder.lock: + if holder.monitor is None or not holder.monitor.alive: + notifier = TelegramNotifier(settings.telegram_token, settings.telegram_chat_id) + holder.monitor = Monitor(detector, store, settings, options, notifier, get_sender()) + holder.monitor.start() + st.rerun() + + st.title("🛡️ EcoWatch AI") + st.caption("Обнаружение пожаров и дыма в реальном времени") + + t_live, t_img, t_an, t_map, t_log = st.tabs(["Мониторинг", "Анализ снимка", "Аналитика", "Карта", "Журнал"]) + with t_live: + live_view(holder) + with t_img: + tab_image(detector, params) + df = incidents_df(store) + with t_an: + tab_analytics(df) + with t_map: + tab_map(df) + with t_log: + tab_journal(df, store) diff --git a/logo/logo_small.png b/logo/logo_small.png new file mode 100644 index 0000000..2918ba2 Binary files /dev/null and b/logo/logo_small.png differ diff --git a/main.py b/main.py index 083a2a1..cbb083d 100644 --- a/main.py +++ b/main.py @@ -1,873 +1,5 @@ -import streamlit as st -import pandas as pd -import plotly.express as px -import plotly.graph_objects as go -from datetime import datetime, timedelta -import cv2 -from ultralytics import YOLO -import tempfile -import numpy as np -import os -from mss import mss -import time -import json -from PIL import Image, ImageDraw, ImageFont -from collections import deque +"""Точка входа веб-интерфейса: `streamlit run main.py`.""" -st.set_page_config(page_title="EcoWatch AI", page_icon="🛡️", layout="wide", initial_sidebar_state="expanded") -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -LOG_FILE = os.path.join(SCRIPT_DIR, "detection_history.json") -RECORD_DIR = os.path.join(SCRIPT_DIR, "recordings") -DANGER_CLASSES = {"fire", "trash", "smoke"} -MIN_CONF = 0.55 -DETECT_CONF = 0.40 -CLASSIFY_CONF = 0.72 -CLASSIFY_GAP = 0.18 -SMOOTH_WINDOW = 5 +from ecowatch.ui import main -_defaults = { - "detection_log": [], "session_start": datetime.now(), "total_frames": 0, - "alert_count": 0, "last_tg_time": datetime.min, "tg_sent_count": 0, - "video_writer": None, "record_start": 0, - "label_buffer": deque(maxlen=SMOOTH_WINDOW), -} -for _k, _v in _defaults.items(): - if _k not in st.session_state: - st.session_state[_k] = _v - - -def load_history(): - if os.path.exists(LOG_FILE): - try: - with open(LOG_FILE, "r", encoding="utf-8") as f: - return json.load(f) - except Exception: - return [] - return [] - - -def save_to_history(entry): - h = load_history() - h.append(entry) - h = h[-2000:] - try: - with open(LOG_FILE, "w", encoding="utf-8") as f: - json.dump(h, f, ensure_ascii=False, default=str) - except Exception: - pass - - -def get_all_data(): - history = load_history() - session_log = st.session_state.get("detection_log", []) - seen = {json.dumps(e, default=str) for e in history} - return history + [e for e in session_log if json.dumps(e, default=str) not in seen] - - -st.markdown(""" - -""", unsafe_allow_html=True) - - -@st.cache_resource -def load_model(): - paths = [ - (os.path.join(SCRIPT_DIR, "models", "best.pt"), "classify"), - (os.path.join(SCRIPT_DIR, "shop project", "runs", "detect", "train", "weights", "best.pt"), "detect"), - (os.path.join(SCRIPT_DIR, "shop project", "runs", "classify", "train_v2", "weights", "best.pt"), "classify"), - (os.path.join(SCRIPT_DIR, "shop project", "runs", "classify", "train", "weights", "best.pt"), "classify"), - ] - for path, mtype in paths: - if os.path.exists(path): - return YOLO(path), mtype - st.error("Модель не найдена. Поместите best.pt в папку models/") - return None, None - - -model = None -MODEL_TYPE = None -try: - model, MODEL_TYPE = load_model() - if MODEL_TYPE == "classify": - st.sidebar.warning("⚠️ Используется classify-модель. Для точного обнаружения обучите detect-модель (train.py)") -except Exception as e: - st.error(f"Ошибка модели: {e}") - -BOT_TOKEN = "8084419721:AAFRMEcaoORZwbCwLpnEKmmmtVK_bG-Xunw" -CHAT_ID = "1719237834" - - -def send_telegram_alert(message, image_array=None): - import requests, io - try: - r = requests.post( - f"https://api.telegram.org/bot{BOT_TOKEN}/sendMessage", - json={"chat_id": CHAT_ID, "text": message, "parse_mode": "HTML"}, - timeout=10, - ) - if not r.json().get("ok"): - return False - if image_array is not None: - ok, buf = cv2.imencode(".jpg", image_array) - if ok: - requests.post( - f"https://api.telegram.org/bot{BOT_TOKEN}/sendPhoto", - data={"chat_id": CHAT_ID, "caption": f"Снимок {datetime.now():%d.%m.%Y %H:%M:%S}"}, - files={"photo": ("alert.jpg", io.BytesIO(buf), "image/jpeg")}, - timeout=10, - ) - st.session_state["tg_sent_count"] += 1 - return True - except Exception: - return False - - -def send_daily_report(): - history = load_history() - today = datetime.now().strftime("%Y-%m-%d") - today_events = [e for e in history if e.get("time", "").startswith(today)] - fires = sum(1 for e in today_events if e.get("label") == "fire") - trash = sum(1 for e in today_events if e.get("label") == "trash") - factory = sum(1 for e in today_events if e.get("label") == "factory") - avg = sum(e.get("confidence", 0) for e in today_events) / max(len(today_events), 1) - msg = ( - f"📊 ЕЖЕДНЕВНЫЙ ОТЧЁТ\n" - f"📅 {datetime.now():%d.%m.%Y}\n\n" - f"🔥 Пожаров: {fires}\n" - f"🗑️ Сжигание мусора: {trash}\n" - f"🏭 Заводской дым: {factory}\n" - f"📈 Средняя вероятность: {avg:.1f}%\n" - f"📋 Всего: {len(today_events)}\n\n" - f"{'⚠️ Повышенная активность!' if fires + trash > 5 else '✅ Ситуация стабильна'}\n\n" - f"— EcoWatch AI" - ) - return send_telegram_alert(msg) - - -def process_frame(results, frame): - r = results[0] - - if hasattr(r, 'boxes') and r.boxes is not None and len(r.boxes) > 0: - boxes = r.boxes - names = r.names - detections = [] - class_counts = {} - best_conf = 0.0 - best_label = "normal" - - for i in range(len(boxes)): - conf = float(boxes.conf[i]) - if conf < DETECT_CONF: - continue - cls_id = int(boxes.cls[i]) - cls_name = names.get(cls_id, f"class_{cls_id}") - x1, y1, x2, y2 = boxes.xyxy[i].cpu().numpy().astype(int) - detections.append((x1, y1, x2, y2, conf, cls_name)) - - class_counts[cls_name] = class_counts.get(cls_name, 0) + 1 - if conf > best_conf: - best_conf = conf - best_label = cls_name - - if not detections: - return "normal", 0.0, {}, [] - - total_det = sum(class_counts.values()) - scores = {name: class_counts.get(name, 0) / total_det for name in names.values()} - return best_label, best_conf, scores, detections - - if hasattr(r, 'probs') and r.probs is not None: - raw = r.probs.data.cpu().numpy().flatten() - scores = {} - for i, name in r.names.items(): - scores[name] = float(raw[i]) - - final_label = max(scores, key=scores.get) - final_conf = scores[final_label] - - if final_conf < CLASSIFY_CONF: - return "normal", final_conf, scores, [] - sorted_vals = sorted(scores.values(), reverse=True) - if len(sorted_vals) >= 2 and (sorted_vals[0] - sorted_vals[1]) < CLASSIFY_GAP: - return "normal", final_conf, scores, [] - if final_label == "normal": - return "normal", final_conf, scores, [] - - return final_label, final_conf, scores, [] - - return "normal", 0.0, {}, [] - - -CLASS_COLORS = { - "fire": (60, 76, 231), - "trash": (11, 160, 232), - "factory": (210, 130, 50), - "smoke": (180, 163, 148), - "normal": (80, 180, 80), -} - -_pil_font_cache = {} - - -def _get_font(size=16): - if size in _pil_font_cache: - return _pil_font_cache[size] - for name in ["arial.ttf", "C:/Windows/Fonts/arial.ttf", "C:/Windows/Fonts/segoeui.ttf"]: - try: - f = ImageFont.truetype(name, size) - _pil_font_cache[size] = f - return f - except Exception: - continue - f = ImageFont.load_default() - _pil_font_cache[size] = f - return f - - -def put_cyrillic(img, text, pos, font_size=16, color=(255, 255, 255)): - pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) - draw = ImageDraw.Draw(pil_img) - font = _get_font(font_size) - draw.text(pos, text, font=font, fill=(color[2], color[1], color[0])) - return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) - - -def get_text_size_pil(text, font_size=16): - font = _get_font(font_size) - bbox = font.getbbox(text) - return bbox[2] - bbox[0], bbox[3] - bbox[1] - - -def smooth_label(label, conf): - buf = st.session_state.get("label_buffer", deque(maxlen=SMOOTH_WINDOW)) - buf.append(label) - st.session_state["label_buffer"] = buf - if label not in DANGER_CLASSES: - return label, conf - if len(buf) >= SMOOTH_WINDOW and all(l in DANGER_CLASSES for l in buf): - return label, conf - return "normal", conf - - -def draw_overlay(frame, label, confidence, scores, detections=None): - overlay = frame.copy() - h, w = frame.shape[:2] - - color = CLASS_COLORS.get(label, (180, 180, 180)) - info = HAZARD_INFO.get(label, {}) - label_ru = info.get("ru", label.upper()) - - if detections: - for (x1, y1, x2, y2, det_conf, cls_name) in detections: - box_color = CLASS_COLORS.get(cls_name, (180, 180, 180)) - det_info = HAZARD_INFO.get(cls_name, {}) - box_label = det_info.get("ru", cls_name.upper()) - - cv2.rectangle(overlay, (x1, y1), (x2, y2), box_color, 2) - - text = f"{box_label} {det_conf*100:.0f}%" - tw, th = get_text_size_pil(text, 14) - cv2.rectangle(overlay, (x1, y1 - th - 10), (x1 + tw + 8, y1), box_color, -1) - overlay = put_cyrillic(overlay, text, (x1 + 4, y1 - th - 6), font_size=14, color=(255, 255, 255)) - - bw, bh = x2 - x1, y2 - y1 - cl = min(20, bw // 4, bh // 4) - ct = 3 - cv2.line(overlay, (x1, y1), (x1 + cl, y1), box_color, ct) - cv2.line(overlay, (x1, y1), (x1, y1 + cl), box_color, ct) - cv2.line(overlay, (x2, y1), (x2 - cl, y1), box_color, ct) - cv2.line(overlay, (x2, y1), (x2, y1 + cl), box_color, ct) - cv2.line(overlay, (x1, y2), (x1 + cl, y2), box_color, ct) - cv2.line(overlay, (x1, y2), (x1, y2 - cl), box_color, ct) - cv2.line(overlay, (x2, y2), (x2 - cl, y2), box_color, ct) - cv2.line(overlay, (x2, y2), (x2, y2 - cl), box_color, ct) - - bar_h = 36 - bar_overlay = overlay.copy() - cv2.rectangle(bar_overlay, (0, 0), (w, bar_h), (0, 0, 0), -1) - cv2.addWeighted(bar_overlay, 0.65, overlay, 0.35, 0, overlay) - - dot_color = (0, 0, 255) if label in DANGER_CLASSES else (0, 200, 0) - cv2.circle(overlay, (18, bar_h // 2), 5, dot_color, -1) - - status_text = f" {label_ru} | {confidence*100:.0f}% | EcoWatch AI" - overlay = put_cyrillic(overlay, status_text, (30, bar_h // 2 - 8), font_size=15, color=(220, 220, 220)) - - time_text = datetime.now().strftime("%H:%M:%S") - ttw, _ = get_text_size_pil(time_text, 14) - overlay = put_cyrillic(overlay, time_text, (w - ttw - 12, bar_h // 2 - 7), font_size=14, color=(150, 150, 150)) - - bottom_h = 28 - bot_y = h - bottom_h - bot_overlay = overlay.copy() - cv2.rectangle(bot_overlay, (0, bot_y), (w, h), (0, 0, 0), -1) - cv2.addWeighted(bot_overlay, 0.6, overlay, 0.4, 0, overlay) - - bar_x = 10 - bar_w_max = w - 20 - bar_y = bot_y + 8 - bar_height = 12 - cv2.rectangle(overlay, (bar_x, bar_y), (bar_x + bar_w_max, bar_y + bar_height), (40, 40, 40), -1) - fill_w = int(bar_w_max * confidence) - cv2.rectangle(overlay, (bar_x, bar_y), (bar_x + fill_w, bar_y + bar_height), color, -1) - - if not detections and label in DANGER_CLASSES and confidence > MIN_CONF: - cv2.rectangle(overlay, (2, bar_h + 2), (w - 2, bot_y - 2), color, 3) - - return overlay - - -def build_probs_html(scores): - sorted_items = sorted(scores.items(), key=lambda x: x[1], reverse=True) - top_name = sorted_items[0][0] - lines = [] - for name, score in sorted_items: - p = score * 100 - info = HAZARD_INFO.get(name, {}) - icon = info.get("icon", "") - ru = info.get("ru", name) - clr = info.get("color", "#8b8fa3") - bar_w = max(int(p), 2) - bold = "font-weight:700;" if name == top_name else "font-weight:400;opacity:.65;" - lines.append( - f"
" - f"{icon} {ru}" - f"{p:.1f}%" - f"
" - f"
" - ) - html = "\n".join(lines) - - if len(sorted_items) >= 2: - gap = (sorted_items[0][1] - sorted_items[1][1]) * 100 - if gap < 12: - html += ( - "
" - "Низкая разница — результат неоднозначен
" - ) - - return html - - -HAZARD_INFO = { - "trash": { - "ru": "СЖИГАНИЕ МУСОРА", "icon": "🗑️", "color": "#eab308", - "desc": "Серый пятнистый дым · Хаотичная форма · Бытовые отходы", - "tg_desc": "• Серый пятнистый дым\n• Хаотичная форма\n• Бытовые отходы, пластик", - }, - "fire": { - "ru": "ПОЖАР", "icon": "🔥", "color": "#ef4444", - "desc": "Густой чёрный/белый дым · Вертикальный столб · Критическая опасность", - "tg_desc": "• Густой чёрный/белый дым\n• Вертикальный столб\n• Критическая опасность", - }, - "factory": { - "ru": "ПРОИЗВОДСТВЕННЫЙ ДЫМ", "icon": "🏭", "color": "#3b82f6", - "desc": "Белый однородный дым · Ровная струя · Промышленный выброс", - "tg_desc": "• Белый однородный дым\n• Ровная струя\n• Промышленный выброс", - }, - "smoke": { - "ru": "ДЫМ", "icon": "💨", "color": "#94a3b8", - "desc": "Обнаружен дым · Возможная угроза · Требуется внимание", - "tg_desc": "• Обнаружен дым\n• Возможная угроза\n• Требуется внимание", - }, - "normal": { - "ru": "НОРМА", "icon": "✅", "color": "#22c55e", - "desc": "Угроз не обнаружено · Чистый кадр", - "tg_desc": "", - }, -} - -with st.sidebar: - logo_path = os.path.join(SCRIPT_DIR, "logo", "logo.png") - if os.path.exists(logo_path): - st.image(logo_path, width=140) - else: - st.markdown( - "
" - "
🛡️
" - "
EcoWatch AI
" - "
v3.1 · Smoke Detection
" - "
", - unsafe_allow_html=True, - ) - - st.markdown('

Источник

', unsafe_allow_html=True) - source_type = st.selectbox( - "Источник", ["Веб-камера", "IP-камера", "Загрузить видео", "Загрузить изображение", "Захват экрана"], - label_visibility="collapsed", - ) - - if source_type == "Веб-камера": - camera_id = st.number_input("ID камеры", 0, 10, 0) - elif source_type == "IP-камера": - ip_url = st.text_input("RTSP URL", "rtsp://") - elif source_type == "Загрузить видео": - uploaded_video = st.file_uploader("Видео", type=["mp4", "avi", "mov"], label_visibility="collapsed") - elif source_type == "Загрузить изображение": - uploaded_image = st.file_uploader("Изображение", type=["jpg", "jpeg", "png"], label_visibility="collapsed") - - st.markdown("---") - st.markdown('

Telegram

', unsafe_allow_html=True) - enable_telegram = st.checkbox("Отправлять алерты", value=True) - send_photo_tg = st.checkbox("Прикреплять фото", value=True) - tg_cooldown = st.slider("Интервал (сек)", 5, 60, 15) - tg_threshold = st.slider("Мин. уверенность (%)", 30, 95, 50) - - c1, c2 = st.columns(2) - with c1: - if st.button("🧪 Тест"): - ok = send_telegram_alert( - "✅ Тест\n\nEcoWatch AI работает.\n" + datetime.now().strftime("%d.%m.%Y %H:%M:%S") - ) - st.success("OK") if ok else st.error("Ошибка") - with c2: - if st.button("📊 Отчёт"): - ok = send_daily_report() - st.success("Отправлен") if ok else st.error("Ошибка") - - st.markdown("---") - st.markdown('

Дополнительно

', unsafe_allow_html=True) - enable_sound = st.checkbox("Звуковой сигнал", value=False) - enable_recording = st.checkbox("Запись при угрозе", value=False) - if enable_recording: - record_duration = st.slider("Длительность записи (сек)", 5, 60, 15) - - st.markdown( - f"
" - f"📨 TG: {st.session_state.get('tg_sent_count', 0)}  " - f"🚨 Алертов: {st.session_state.get('alert_count', 0)}
" - f"🎞️ Кадров: {st.session_state.get('total_frames', 0)}  " - f"⏱️ Сессия: {str(datetime.now() - st.session_state['session_start']).split('.')[0]}" - f"
", - unsafe_allow_html=True, - ) - -st.markdown("

🛡️ EcoWatch AI

", unsafe_allow_html=True) -st.markdown( - "

" - "Система обнаружения экологических угроз

", - unsafe_allow_html=True, -) - -tab1, tab2, tab3, tab4 = st.tabs(["Мониторинг", "Аналитика", "Карта", "Журнал"]) - -with tab1: - col_v, col_s = st.columns([5, 2]) - with col_v: - video_placeholder = st.empty() - is_running = st.toggle("Запустить мониторинг", value=False, key="run_toggle") - if not is_running: - video_placeholder.markdown( - "
" - "
" - "
📹
" - "

Ожидание запуска

" - "
", - unsafe_allow_html=True, - ) - with col_s: - alert_placeholder = st.empty() - stats_placeholder = st.empty() - sound_placeholder = st.empty() - if not is_running: - alert_placeholder.markdown( - "
" - "
Система готова
" - "
" - "Включите мониторинг для анализа
", - unsafe_allow_html=True, - ) - - if is_running and model: - cap = None - sct = None - frame_count = 0 - start_time = time.time() - - try: - if source_type == "Веб-камера": - cap = cv2.VideoCapture(int(camera_id)) - elif source_type == "IP-камера": - cap = cv2.VideoCapture(ip_url) - elif source_type == "Загрузить видео" and uploaded_video: - tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") - tfile.write(uploaded_video.read()) - cap = cv2.VideoCapture(tfile.name) - elif source_type == "Захват экрана": - sct = mss() - monitor = {"top": 100, "left": 100, "width": 800, "height": 600} - elif source_type == "Загрузить изображение" and uploaded_image: - file_bytes = np.asarray(bytearray(uploaded_image.read()), dtype=np.uint8) - frame = cv2.imdecode(file_bytes, 1) - results = model(frame) - - label, conf, scores, detections = process_frame(results, frame) - info = HAZARD_INFO.get(label, {}) - label_ru = info.get("ru", label.upper()) - is_img_danger = label in DANGER_CLASSES and conf > MIN_CONF - probs_html = build_probs_html(scores) - - annotated = draw_overlay(frame, label, conf, scores, detections) - video_placeholder.image(annotated, channels="BGR") - - if is_img_danger: - alert_placeholder.markdown( - f"
" - f"
{info.get('icon', '')} {label_ru}
" - f"
{info.get('desc', '')}
" - f"
{probs_html}
" - f"
", - unsafe_allow_html=True, - ) - entry = { - "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), - "type": label_ru, "label": label, - "confidence": round(conf * 100, 1), "source": source_type, - } - st.session_state["detection_log"].append(entry) - save_to_history(entry) - st.session_state["alert_count"] += 1 - else: - alert_placeholder.markdown( - f"
" - f"
Угроз не обнаружено
" - f"
{probs_html}
" - f"
", - unsafe_allow_html=True, - ) - except Exception as e: - st.error(f"Ошибка: {e}") - - if cap or sct: - while is_running: - if sct: - img = np.array(sct.grab(monitor)) - frame = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) - else: - ret, frame = cap.read() - if not ret: - st.warning("Конец потока") - break - - results = model(frame) - label, conf, scores, detections = process_frame(results, frame) - label, conf = smooth_label(label, conf) - - annotated = draw_overlay(frame, label, conf, scores, detections) - video_placeholder.image(annotated, channels="BGR") - - frame_count += 1 - st.session_state["total_frames"] += 1 - elapsed = time.time() - start_time - fps = frame_count / elapsed if elapsed > 0 else 0 - - info = HAZARD_INFO.get(label, {}) - label_ru = info.get("ru", label.upper()) - is_danger = label in DANGER_CLASSES and conf > MIN_CONF - probs_html = build_probs_html(scores) - - if is_danger: - alert_placeholder.markdown( - f"
" - f"
{info.get('icon', '')} {label_ru}
" - f"
{info.get('desc', '')}
" - f"
{probs_html}
" - f"
{datetime.now():%H:%M:%S}
" - f"
", - unsafe_allow_html=True, - ) - if enable_sound: - sound_placeholder.markdown( - '', - unsafe_allow_html=True, - ) - else: - alert_placeholder.markdown( - f"
" - f"
Норма
" - f"
{probs_html}
" - f"
{datetime.now():%H:%M:%S}
", - unsafe_allow_html=True, - ) - sound_placeholder.empty() - - dur = str(timedelta(seconds=int(elapsed))) - stats_placeholder.markdown( - f"
" - f"
{fps:.0f}
FPS
" - f"
{frame_count}
Кадров
" - f"
{st.session_state['alert_count']}
Алертов
" - f"
{dur}
Время
" - f"
", - unsafe_allow_html=True, - ) - - if is_danger and frame_count % 30 == 0: - entry = { - "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), - "type": label_ru, "label": label, - "confidence": round(conf * 100, 1), "source": source_type, - } - st.session_state["detection_log"].append(entry) - save_to_history(entry) - st.session_state["alert_count"] += 1 - - if enable_telegram and is_danger and conf > tg_threshold / 100.0: - diff = (datetime.now() - st.session_state.get("last_tg_time", datetime.min)).total_seconds() - if diff > tg_cooldown: - msg = ( - f"🚨 ЭКОЛОГИЧЕСКАЯ УГРОЗА\n\n" - f"{info.get('icon', '')} {label_ru}\n" - f"📊 Вероятность: {conf * 100:.1f}%\n" - f"📹 Источник: {source_type}\n" - f"🕐 {datetime.now():%d.%m.%Y %H:%M:%S}\n\n" - f"{info.get('tg_desc', '')}\n\n" - f"— EcoWatch AI" - ) - send_telegram_alert(msg, annotated if send_photo_tg else None) - st.session_state["last_tg_time"] = datetime.now() - - if enable_recording and is_danger: - if st.session_state.get("video_writer") is None: - os.makedirs(RECORD_DIR, exist_ok=True) - fname = os.path.join(RECORD_DIR, f"alert_{datetime.now():%Y%m%d_%H%M%S}.avi") - fh, fw = frame.shape[:2] - st.session_state["video_writer"] = cv2.VideoWriter( - fname, cv2.VideoWriter_fourcc(*"XVID"), 15, (fw, fh) - ) - st.session_state["record_start"] = time.time() - wr = st.session_state.get("video_writer") - if wr: - wr.write(annotated) - if time.time() - st.session_state.get("record_start", 0) > record_duration: - wr.release() - st.session_state["video_writer"] = None - - if cap: - cap.release() - wr = st.session_state.get("video_writer") - if wr: - wr.release() - st.session_state["video_writer"] = None - - -_PLOT = dict( - plot_bgcolor="rgba(0,0,0,0)", paper_bgcolor="rgba(0,0,0,0)", - font=dict(color="#8b8fa3", family="Inter"), - margin=dict(t=30, b=20, l=20, r=20), - xaxis=dict(gridcolor="rgba(255,255,255,.04)", zerolinecolor="rgba(255,255,255,.04)"), - yaxis=dict(gridcolor="rgba(255,255,255,.04)", zerolinecolor="rgba(255,255,255,.04)"), -) - -with tab2: - all_data = get_all_data() - if len(all_data) > 0: - df = pd.DataFrame(all_data) - df["time"] = pd.to_datetime(df["time"]) - - m1, m2, m3, m4 = st.columns(4) - with m1: - st.metric("Всего", len(df)) - with m2: - st.metric("Пожаров", len(df[df["label"] == "fire"])) - with m3: - st.metric("Мусор", len(df[df["label"] == "trash"])) - with m4: - st.metric("Заводской", len(df[df["label"] == "factory"])) - - c1, c2 = st.columns(2) - with c1: - daily = df.groupby(df["time"].dt.date).size().reset_index(name="count") - fig = px.bar(daily, x="time", y="count", labels={"time": "", "count": ""}, color_discrete_sequence=["#22c55e"]) - fig.update_layout(**_PLOT, title="События по дням") - st.plotly_chart(fig, use_container_width=True) - - with c2: - tc = df["type"].value_counts().reset_index() - tc.columns = ["Тип", "N"] - fig = px.pie(tc, values="N", names="Тип", color_discrete_sequence=["#ef4444", "#eab308", "#3b82f6"], hole=0.45) - fig.update_layout( - plot_bgcolor="rgba(0,0,0,0)", paper_bgcolor="rgba(0,0,0,0)", - font=dict(color="#8b8fa3", family="Inter"), margin=dict(t=30, b=20), title="Распределение", - ) - fig.update_traces(textinfo="percent+label", textfont_size=11) - st.plotly_chart(fig, use_container_width=True) - - c3, c4 = st.columns(2) - with c3: - fig = px.scatter( - df.tail(100), x="time", y="confidence", color="type", size="confidence", - labels={"time": "", "confidence": "%", "type": "Тип"}, - color_discrete_sequence=["#ef4444", "#eab308", "#3b82f6"], - ) - fig.update_layout(**_PLOT, title="Вероятности обнаружений", showlegend=True, legend=dict(orientation="h", y=-0.15)) - st.plotly_chart(fig, use_container_width=True) - - with c4: - avg = df.groupby("type")["confidence"].mean().reset_index() - avg.columns = ["Тип", "Ср. %"] - fig = px.bar(avg, x="Тип", y="Ср. %", color="Тип", color_discrete_sequence=["#ef4444", "#eab308", "#3b82f6"]) - fig.update_layout(**_PLOT, title="Средняя уверенность", showlegend=False) - st.plotly_chart(fig, use_container_width=True) - - if len(df) > 5: - st.markdown("#### Активность по часам") - df["hour"] = df["time"].dt.hour - df["weekday"] = df["time"].dt.day_name() - day_order = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"] - heat = df.groupby(["weekday", "hour"]).size().reset_index(name="count") - hp = heat.pivot_table(index="weekday", columns="hour", values="count", fill_value=0) - hp = hp.reindex([d for d in day_order if d in hp.index]) - day_ru = {"Monday": "Пн", "Tuesday": "Вт", "Wednesday": "Ср", "Thursday": "Чт", "Friday": "Пт", "Saturday": "Сб", "Sunday": "Вс"} - fig = go.Figure( - data=go.Heatmap( - z=hp.values, - x=[f"{hh}:00" for hh in hp.columns], - y=[day_ru.get(d, d) for d in hp.index], - colorscale=[[0, "#0f1117"], [0.5, "#22c55e"], [1, "#ef4444"]], - hovertemplate="%{y} %{x}
Событий: %{z}", - ) - ) - fig.update_layout( - plot_bgcolor="rgba(0,0,0,0)", paper_bgcolor="rgba(0,0,0,0)", - font=dict(color="#8b8fa3", family="Inter"), margin=dict(t=10, b=10, l=10, r=10), height=220, - ) - st.plotly_chart(fig, use_container_width=True) - else: - st.markdown( - "
" - "
📊
" - "

Данные появятся после запуска мониторинга

", - unsafe_allow_html=True, - ) - -with tab3: - st.markdown("#### Карта обнаружений") - st.caption("Координаты привязаны к расположению камер.") - - all_data_map = get_all_data() - base_lat = st.number_input("Широта камеры", value=51.1694, format="%.4f", key="map_lat") - base_lon = st.number_input("Долгота камеры", value=71.4491, format="%.4f", key="map_lon") - - if len(all_data_map) > 0: - map_points = [] - for i, e in enumerate(all_data_map[-100:]): - np.random.seed(hash(e.get("time", "")) % 2**31) - map_points.append({"lat": base_lat + np.random.uniform(-0.015, 0.015), "lon": base_lon + np.random.uniform(-0.015, 0.015)}) - df_map = pd.DataFrame(map_points) - st.map(df_map, latitude="lat", longitude="lon") - st.markdown( - f"
" - f"Показано {len(map_points)}" - f" последних обнаружений
", - unsafe_allow_html=True, - ) - else: - st.markdown( - "
" - "
🗺️
" - "

Нет данных

", - unsafe_allow_html=True, - ) - -with tab4: - all_data_log = get_all_data() - if len(all_data_log) > 0: - df_l = pd.DataFrame(all_data_log) - df_l["time"] = pd.to_datetime(df_l["time"]) - df_l = df_l.sort_values("time", ascending=False) - - fc1, fc2 = st.columns(2) - with fc1: - f_type = st.multiselect("Тип", df_l["type"].unique(), default=list(df_l["type"].unique())) - with fc2: - f_conf = st.slider("Мин. уверенность", 0, 100, 0, key="log_conf") - - filt = df_l[(df_l["type"].isin(f_type)) & (df_l["confidence"] >= f_conf)] - disp = filt.rename(columns={"time": "Время", "type": "Тип", "confidence": "%", "source": "Источник"}) - cols = ["Время", "Тип", "%", "Источник"] - disp = disp[[c for c in cols if c in disp.columns]] - st.dataframe(disp, use_container_width=True, height=400) - - bc1, bc2, bc3 = st.columns(3) - with bc1: - csv = disp.to_csv(index=False).encode("utf-8") - st.download_button("Скачать CSV", csv, f"ecowatch_{datetime.now():%Y%m%d}.csv", "text/csv") - with bc2: - report = ( - f"ОТЧЁТ ECOWATCH AI\n{'=' * 40}\n" - f"Дата: {datetime.now():%d.%m.%Y %H:%M}\n" - f"Событий: {len(filt)}\n" - f"Пожаров: {len(filt[filt['type'].str.contains('ПОЖАР', na=False)])}\n" - f"Мусор: {len(filt[filt['type'].str.contains('МУСОР', na=False)])}\n" - f"{'=' * 40}\n" - ) - for _, r in filt.head(50).iterrows(): - report += f"\n{r.get('time', '')} | {r.get('type', '')} | {r.get('confidence', '')}%" - st.download_button("Скачать отчёт", report.encode("utf-8"), f"report_{datetime.now():%Y%m%d}.txt") - with bc3: - if st.button("Очистить историю"): - if os.path.exists(LOG_FILE): - os.remove(LOG_FILE) - st.session_state["detection_log"] = [] - st.session_state["alert_count"] = 0 - st.rerun() - else: - st.markdown( - "
" - "
📋
" - "

Журнал пуст

", - unsafe_allow_html=True, - ) - -st.markdown( - "
EcoWatch AI · 2025 · Infomatrix
YOLOv8 · Streamlit · OpenCV
", - unsafe_allow_html=True, -) +main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2299489 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "ecowatch" +version = "4.0.0" +description = "EcoWatch AI — обнаружение пожаров и дыма на видео (YOLO + Streamlit)" +requires-python = ">=3.10" + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "W", "I", "B", "UP"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..9d7103c --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,2 @@ +pytest>=8 +ruff>=0.6 diff --git a/requirements.txt b/requirements.txt index 0989f88..f143192 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,9 @@ -streamlit -pandas -plotly -opencv-python -ultralytics -mss -Pillow -requests -numpy \ No newline at end of file +ultralytics>=8.1 +streamlit>=1.48 +pandas>=2.0 +plotly>=5.24 +opencv-python>=4.8 +numpy>=1.24 +Pillow>=10.1 +requests>=2.31 +mss>=9.0 diff --git a/run_app.bat b/run_app.bat index 88c95ed..f345da4 100644 --- a/run_app.bat +++ b/run_app.bat @@ -1,8 +1,7 @@ -@echo off -cd /d "%~dp0" -echo Starting EcoWatch AI... -echo. -echo If the browser does not open automatically, go to: http://localhost:8501 -echo. -"C:\Program Files\Python311\python.exe" -m streamlit run main.py -pause +@echo off +cd /d "%~dp0" +echo Starting EcoWatch AI... +echo If the browser does not open automatically, go to: http://localhost:8501 +echo. +where py >nul 2>nul && (py -3 -m streamlit run main.py) || (python -m streamlit run main.py) +pause diff --git a/start_ecowatch.bat b/start_ecowatch.bat deleted file mode 100644 index 842ff63..0000000 --- a/start_ecowatch.bat +++ /dev/null @@ -1,4 +0,0 @@ -@echo off -echo Starting EcoWatch AI... -"C:\Program Files\Python311\python.exe" -m streamlit run main.py -pause diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ece8e0c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,4 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) diff --git a/tests/test_config_sources.py b/tests/test_config_sources.py new file mode 100644 index 0000000..124403c --- /dev/null +++ b/tests/test_config_sources.py @@ -0,0 +1,29 @@ +from ecowatch.config import Settings +from ecowatch.hazards import get_hazard, is_danger +from ecowatch.sources import parse_source + + +def test_settings_from_dotenv(tmp_path, monkeypatch): + for k in ("ECOWATCH_TG_TOKEN", "ECOWATCH_TG_CHAT_ID", "ECOWATCH_CAMERA_LAT", "ECOWATCH_DATA_DIR"): + monkeypatch.delenv(k, raising=False) + env = tmp_path / ".env" + env.write_text( + '# comment\nECOWATCH_TG_TOKEN="abc"\nexport ECOWATCH_TG_CHAT_ID=42\nECOWATCH_CAMERA_LAT=bad\n' + f"ECOWATCH_DATA_DIR={tmp_path / 'd'}\n" + ) + s = Settings.from_env(env) + assert s.telegram_token == "abc" and s.telegram_chat_id == "42" and s.telegram_configured + assert s.camera_lat is None and s.db_path == tmp_path / "d" / "ecowatch.db" + + +def test_parse_source(): + assert parse_source("0").kind == "webcam" + assert parse_source("screen:2").value == 2 + s = parse_source("rtsp://u:p@10.0.0.1/live") + assert s.kind == "stream" and "p@" not in s.label # пароль не попадает в подписи + assert parse_source("clip.mp4").kind == "file" + + +def test_hazards(): + assert is_danger("Fire") and is_danger("flame") and not is_danger("normal") + assert get_hazard("unknown_cls").title == "Unknown_cls" diff --git a/tests/test_detector.py b/tests/test_detector.py new file mode 100644 index 0000000..22d7f3c --- /dev/null +++ b/tests/test_detector.py @@ -0,0 +1,47 @@ +from types import SimpleNamespace + +import numpy as np + +from ecowatch.detector import parse_result + + +def test_parse_classify_result(): + r = SimpleNamespace( + names={0: "fire", 1: "normal", 2: "Smoke"}, probs=SimpleNamespace(data=np.array([0.7, 0.2, 0.1])) + ) + res = parse_result(r) + assert res.scores == {"fire": 0.7, "normal": 0.2, "smoke": 0.1} + assert res.top() == ("fire", 0.7) + assert res.boxes == [] + + +def test_parse_detect_result_filters_boxes_and_derives_normal(): + boxes = SimpleNamespace( + conf=np.array([0.8, 0.2, 0.5]), + cls=np.array([0, 0, 1]), + xyxy=np.array([[1, 2, 30, 40], [0, 0, 5, 5], [10, 10, 50, 60]], dtype=float), + ) + boxes.__len__ = lambda: 3 # noqa: E731 + r = SimpleNamespace(names={0: "fire", 1: "smoke"}, probs=None, boxes=_Sized(boxes, 3)) + res = parse_result(r, box_conf=0.3) + assert len(res.boxes) == 2 + assert res.scores["fire"] == 0.8 and res.scores["smoke"] == 0.5 + assert abs(res.scores["normal"] - 0.2) < 1e-9 + + +def test_parse_detect_result_empty(): + r = SimpleNamespace(names={0: "fire"}, probs=None, boxes=None) + res = parse_result(r) + assert res.scores == {"fire": 0.0, "normal": 1.0} + assert res.top()[0] == "normal" + + +class _Sized: + def __init__(self, inner, n): + self._inner, self._n = inner, n + + def __len__(self): + return self._n + + def __getattr__(self, item): + return getattr(self._inner, item) diff --git a/tests/test_notifier.py b/tests/test_notifier.py new file mode 100644 index 0000000..d79ed87 --- /dev/null +++ b/tests/test_notifier.py @@ -0,0 +1,69 @@ +from datetime import date, datetime + +import numpy as np + +from ecowatch.notifier import AsyncSender, TelegramNotifier, format_alert, format_daily_report + + +class FakeResp: + def __init__(self, ok=True): + self._ok = ok + self.text = "" + + def json(self): + return {"ok": self._ok} + + +class FakeSession: + def __init__(self, ok=True): + self.calls = [] + self.ok = ok + + def post(self, url, **kw): + self.calls.append((url, kw)) + return FakeResp(self.ok) + + +def test_disabled_without_credentials(): + s = FakeSession() + n = TelegramNotifier(None, "1", session=s) + assert not n.enabled and not n.send_message("hi") and s.calls == [] + + +def test_send_message_and_photo(): + s = FakeSession() + n = TelegramNotifier("T", "42", session=s) + assert n.send_message("hello") + assert n.send_photo(np.zeros((10, 10, 3), np.uint8), "cap") + assert s.calls[0][0].endswith("/botT/sendMessage") + assert s.calls[1][0].endswith("/sendPhoto") and s.calls[1][1]["data"]["chat_id"] == "42" + + +def test_long_caption_split(): + s = FakeSession() + TelegramNotifier("T", "1", session=s).send_photo(np.zeros((4, 4, 3), np.uint8), "x" * 2000) + assert [c[0].rsplit("/", 1)[1] for c in s.calls] == ["sendMessage", "sendPhoto"] + + +def test_async_sender_counts(): + sender = AsyncSender() + sender.submit(lambda: True) + sender.submit(lambda: False) + sender.submit(lambda: 1 / 0) + sender.join() + assert sender.sent == 1 and sender.failed == 2 + + +def test_format_alert_escapes_html(): + text = format_alert("fire", 0.876, "", datetime(2026, 1, 1, 9, 5)) + assert "ПОЖАР" in text and "88%" in text and "<cam>" in text and "01.01.2026 09:05:00" in text + + +def test_daily_report(): + text = format_daily_report( + {"day": date(2026, 1, 1), "total": 3, "by_label": {"fire": 2, "smoke": 1}, "peak_conf": 0.9} + ) + assert "Пожар: 2" in text and "Всего инцидентов: 3" in text and "стабильна" in text + assert "не зафиксировано" in format_daily_report( + {"day": date(2026, 1, 1), "total": 0, "by_label": {}, "peak_conf": 0} + ) diff --git a/tests/test_recorder.py b/tests/test_recorder.py new file mode 100644 index 0000000..9554cc5 --- /dev/null +++ b/tests/test_recorder.py @@ -0,0 +1,37 @@ +import cv2 +import numpy as np + +from ecowatch.recorder import ClipRecorder + + +def frame(v): + return np.full((48, 64, 3), v, np.uint8) + + +def count_frames(path): + cap = cv2.VideoCapture(str(path)) + n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + return n + + +def test_clip_has_preroll_and_postroll(tmp_path): + rec = ClipRecorder(tmp_path, fps=10, pre_seconds=1, post_seconds=0.55) + t = 0.0 + for _ in range(30): # 3 с без тревоги — в клип попадёт только последняя секунда + assert rec.push(frame(0), False, t) is None + t += 0.1 + for _ in range(5): + rec.push(frame(255), True, t) + t += 0.1 + assert rec.recording + path = None + while path is None: + path = rec.push(frame(0), False, t) + t += 0.1 + assert path.exists() and not rec.recording + assert count_frames(path) == 10 + 5 + 6 + + +def test_close_without_recording(tmp_path): + assert ClipRecorder(tmp_path).close() is None diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..b28cdc6 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,68 @@ +import json +from datetime import date, datetime + +import pytest + +from ecowatch.storage import EventStore + + +@pytest.fixture +def store(tmp_path): + s = EventStore(tmp_path / "db.sqlite") + yield s + s.close() + + +def test_open_update_list(store): + i = store.open_incident("fire", 0.8, "Камера 0", started_at=datetime(2026, 1, 2, 10, 0, 0), lat=51.1, lon=71.4) + store.update_incident(i, ended_at=datetime(2026, 1, 2, 10, 1, 0), peak_conf=0.93, clip="/x.mp4") + (row,) = store.list() + assert row["label"] == "fire" and row["peak_conf"] == 0.93 + assert row["ended_at"] == "2026-01-02T10:01:00" and row["clip"] == "/x.mp4" and row["lat"] == 51.1 + + +def test_update_rejects_unknown_fields(store): + i = store.open_incident("smoke", 0.7, "s") + with pytest.raises(ValueError): + store.update_incident(i, id=5) + + +def test_list_order_since_limit_and_clear(store): + for h in (8, 9, 10): + store.open_incident("smoke", 0.7, "s", started_at=datetime(2026, 1, 2, h)) + assert [r["started_at"][11:13] for r in store.list()] == ["10", "09", "08"] + assert len(store.list(since=datetime(2026, 1, 2, 9))) == 2 + assert len(store.list(limit=1)) == 1 + store.clear() + assert store.count() == 0 + + +def test_day_summary(store): + store.open_incident("fire", 0.9, "s", started_at=datetime(2026, 3, 1, 12)) + store.open_incident("fire", 0.7, "s", started_at=datetime(2026, 3, 1, 13)) + store.open_incident("smoke", 0.6, "s", started_at=datetime(2026, 3, 1, 14)) + store.open_incident("fire", 0.99, "s", started_at=datetime(2026, 3, 2, 1)) + s = store.day_summary(date(2026, 3, 1)) + assert s["total"] == 3 and s["by_label"] == {"fire": 2, "smoke": 1} and s["peak_conf"] == 0.9 + + +def test_import_legacy_json(store, tmp_path): + legacy = tmp_path / "detection_history.json" + legacy.write_text( + json.dumps( + [ + { + "time": "2025-05-01 12:00:00", + "type": "ПОЖАР", + "label": "fire", + "confidence": 87.5, + "source": "Веб-камера", + }, + {"time": "broken"}, + ] + ), + encoding="utf-8", + ) + assert store.import_legacy_json(legacy) == 1 + (row,) = store.list() + assert row["label"] == "fire" and row["peak_conf"] == pytest.approx(0.875) diff --git a/tests/test_tracker.py b/tests/test_tracker.py new file mode 100644 index 0000000..16c94d0 --- /dev/null +++ b/tests/test_tracker.py @@ -0,0 +1,59 @@ +from ecowatch.config import DetectionParams +from ecowatch.detector import FrameResult +from ecowatch.tracker import ThreatTracker, classify_frame, evaluate_single + +FIRE = FrameResult({"fire": 0.9, "normal": 0.05, "smoke": 0.05}) +CLEAN = FrameResult({"fire": 0.02, "normal": 0.95, "smoke": 0.03}) +PARAMS = DetectionParams(conf_threshold=0.6, margin=0.15, ema_alpha=1.0, confirm_frames=3, clear_frames=2) + + +def test_classify_frame_requires_threshold_and_margin(): + p = DetectionParams(conf_threshold=0.6, margin=0.2) + assert classify_frame({"fire": 0.9, "normal": 0.1}, p)[2] + assert not classify_frame({"fire": 0.55, "normal": 0.45}, p)[2] # ниже порога + assert not classify_frame({"fire": 0.62, "smoke": 0.5}, p)[2] # мал отрыв + assert not classify_frame({"normal": 0.9, "fire": 0.1}, p)[2] + assert not classify_frame({}, p)[2] + + +def test_onset_after_confirm_frames_and_clear_after_clear_frames(): + t = ThreatTracker(PARAMS) + decisions = [t.update(FIRE) for _ in range(3)] + assert [d.active for d in decisions] == [False, False, True] + assert decisions[0].frame_danger and decisions[0].label == "normal" + assert decisions[2].onset and decisions[2].label == "fire" + assert sum(d.onset for d in decisions + [t.update(FIRE) for _ in range(5)]) == 1 + + d1 = t.update(CLEAN) + assert d1.active and not d1.cleared # одного чистого кадра мало + d2 = t.update(CLEAN) + assert d2.cleared and not d2.active and d2.label == "normal" + + +def test_single_noisy_frame_does_not_trigger(): + t = ThreatTracker(PARAMS) + seq = [CLEAN, FIRE, CLEAN, FIRE, FIRE, CLEAN] + assert not any(t.update(r).active for r in seq) + + +def test_ema_suppresses_spike(): + t = ThreatTracker(DetectionParams(ema_alpha=0.3, confirm_frames=1)) + for _ in range(10): + t.update(CLEAN) + assert not t.update(FIRE).frame_danger # один всплеск не перебивает историю + + +def test_escalation_from_smoke_to_fire(): + t = ThreatTracker(PARAMS) + smoke = FrameResult({"smoke": 0.9, "normal": 0.05, "fire": 0.05}) + for _ in range(3): + d = t.update(smoke) + assert d.label == "smoke" + d = t.update(FIRE) + assert d.active and d.label == "fire" and not d.onset + + +def test_evaluate_single(): + assert evaluate_single(FIRE).label == "fire" + d = evaluate_single(CLEAN) + assert d.label == "normal" and d.confidence == 0.95 and not d.active diff --git a/train.py b/train.py index 52ff411..ae574b5 100644 --- a/train.py +++ b/train.py @@ -1,94 +1,114 @@ -import os -import sys - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +"""Обучение модели EcoWatch AI. -DATASET_DIR = os.path.join(SCRIPT_DIR, "dataset_detect") -PROJECT_DIR = os.path.join(SCRIPT_DIR, "shop project", "runs", "classify") +Классификация (как текущая models/best.pt) — датасет из папок по классам: + python train.py --data dataset_detect -BASE_MODEL = "yolov8s-cls.pt" -EPOCHS = 50 -IMG_SIZE = 224 -BATCH_SIZE = 32 -PATIENCE = 15 +Детекция с рамками — датасет в формате YOLO с data.yaml: + python train.py --data fire_smoke/data.yaml --model yolov8s.pt --imgsz 640 +Лучшие веса копируются в models/best.pt, прежняя модель сохраняется как models/best_<дата>.pt. +""" -def check_dataset(): - train_dir = os.path.join(DATASET_DIR, "train") - valid_dir = os.path.join(DATASET_DIR, "valid") +from __future__ import annotations - if not os.path.isdir(train_dir): - print(f"Папка train не найдена: {train_dir}") - return False - if not os.path.isdir(valid_dir): - print(f"Папка valid не найдена: {valid_dir}") +import argparse +import shutil +import sys +from datetime import datetime +from pathlib import Path + +ROOT = Path(__file__).resolve().parent +IMAGE_EXT = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} + + +def describe_classify_dataset(data: Path) -> bool: + ok = True + for split in ("train", "val", "valid"): + d = data / split + if not d.is_dir(): + continue + counts = { + c.name: sum(1 for f in c.iterdir() if f.suffix.lower() in IMAGE_EXT) + for c in sorted(d.iterdir()) + if c.is_dir() + } + print(f" {split}: " + ", ".join(f"{k}={v}" for k, v in counts.items()) + f" (всего {sum(counts.values())})") + empty = [k for k, v in counts.items() if v == 0] + if empty: + print(f" ⚠ пустые классы в {split}: {empty}") + ok = False + if not (data / "train").is_dir(): + print(f"Нет папки {data / 'train'}") return False - - print("Датасет:") - total = 0 - for split_name, split_dir in [("train", train_dir), ("valid", valid_dir)]: - classes = sorted(os.listdir(split_dir)) - counts = {} - for cls in classes: - cls_path = os.path.join(split_dir, cls) - if os.path.isdir(cls_path): - n = len([f for f in os.listdir(cls_path) if os.path.isfile(os.path.join(cls_path, f))]) - counts[cls] = n - total += n - parts = ", ".join(f"{k}: {v}" for k, v in counts.items()) - print(f" {split_name}: {parts}") - - print(f" Всего: {total} изображений") - print(f" Классы: {classes}") - return True - - -def train(): - if not check_dataset(): - return + return ok + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--data", default=str(ROOT / "dataset_detect"), help="папка датасета или data.yaml") + p.add_argument("--model", help="базовые веса (по умолчанию yolov8s-cls.pt / yolov8s.pt)") + p.add_argument("--epochs", type=int, default=50) + p.add_argument("--imgsz", type=int, help="размер входа (224 для classify, 640 для detect)") + p.add_argument("--batch", type=int, default=32) + p.add_argument("--patience", type=int, default=15) + p.add_argument("--device", help="cpu, 0, 0,1 … (по умолчанию — автоматически)") + p.add_argument("--name", default=f"run_{datetime.now():%Y%m%d_%H%M}") + p.add_argument("--no-install", action="store_true", help="не копировать результат в models/best.pt") + return p.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + data = Path(args.data) + task = "detect" if data.suffix in {".yaml", ".yml"} else "classify" + print(f"Задача: {task}, датасет: {data}") + if not data.exists(): + print("Датасет не найден") + return 1 + if task == "classify" and not describe_classify_dataset(data): + return 1 from ultralytics import YOLO - print(f"Базовая модель: {BASE_MODEL}") - print(f"Эпох: {EPOCHS}, Размер: {IMG_SIZE}px, Батч: {BATCH_SIZE}") - print(f"Сохранение: {PROJECT_DIR}") - - model = YOLO(BASE_MODEL) - - results = model.train( - data=DATASET_DIR, - epochs=EPOCHS, - imgsz=IMG_SIZE, - batch=BATCH_SIZE, - patience=PATIENCE, - project=PROJECT_DIR, - name="train_v2", + base = args.model or ("yolov8s-cls.pt" if task == "classify" else "yolov8s.pt") + imgsz = args.imgsz or (224 if task == "classify" else 640) + model = YOLO(base) + model.train( + data=str(data), + epochs=args.epochs, + imgsz=imgsz, + batch=args.batch, + patience=args.patience, + device=args.device, + project=str(ROOT / "runs" / task), + name=args.name, exist_ok=True, - pretrained=True, - optimizer="auto", - lr0=0.01, - lrf=0.01, - degrees=10.0, - flipud=0.0, + # Огонь и дым не бывают вверх ногами, а цвет важен — аугментации умеренные. fliplr=0.5, - hsv_h=0.015, - hsv_s=0.7, + flipud=0.0, + degrees=10.0, + hsv_h=0.01, + hsv_s=0.5, hsv_v=0.4, - save=True, plots=True, - verbose=True, ) - best_path = os.path.join(PROJECT_DIR, "train_v2", "weights", "best.pt") - models_dir = os.path.join(SCRIPT_DIR, "models") - if os.path.exists(best_path): - os.makedirs(models_dir, exist_ok=True) - import shutil - shutil.copy2(best_path, os.path.join(models_dir, "best.pt")) - print(f"Модель обучена и скопирована в models/best.pt") - else: - print("Ошибка обучения. Проверь логи выше.") + best = Path(model.trainer.best) + if not best.is_file(): + print("Обучение не дало весов — смотрите лог выше") + return 1 + print(f"Лучшие веса: {best}") + if not args.no_install: + target = ROOT / "models" / "best.pt" + target.parent.mkdir(exist_ok=True) + if target.exists(): + backup = target.with_name(f"best_{datetime.now():%Y%m%d_%H%M%S}.pt") + shutil.copy2(target, backup) + print(f"Старая модель сохранена: {backup.name}") + shutil.copy2(best, target) + print(f"Модель установлена: {target}") + return 0 if __name__ == "__main__": - train() + sys.exit(main())