change image model

readme.md

@@ -3,7 +3,7 @@
 Сервис автоматической модерации пользовательского контента на базе **FastAPI** и моделей машинного обучения. Позволяет проверять тексты и изображения на наличие недопустимого контента:
 
 - **Тексты** — токсичность (нейросетевая модель BERT) и нецензурная лексика (словарь + морфологический анализ `pymorphy3`).
-- **Изображения** — NSFW-контент (нейросетевая модель `Falconsai/nsfw_image_detection`) и запрещённые категории (марихуана, наркотики, оружие, порнография, насилие) через мультимодальную модель CLIP.
+- **Изображения** — NSFW-контент (нейросетевая модель `Falconsai/nsfw_image_detection`), оружие (zero-shot детектор **OWLv2**) и запрещённые категории (марихуана, наркотики, порнография, насилие) через мультимодальную модель CLIP.
 
 Сервис предоставляет REST API и возвращает вердикт (`approved`), оценку уверенности (`score`), причину отказа (`reason`) и дополнительную диагностическую информацию.
 
@@ -39,7 +39,8 @@
 - **Модерация изображений**:
   - Проверка формата (JPEG/PNG/WebP) и размера файла (до 5 МБ);
   - Валидация геометрии изображения (50×50 … 10000×10000 пикселей);
-  - Двухэтапная проверка: модель NSFW-классификации + CLIP-классификация по заранее заданным категориям;
+  - Комбинированная проверка: модель NSFW-классификации + CLIP-классификация по запрещённым категориям (наркотики, насилие);
+  - Отдельный детектор оружия на базе **OWLv2** (zero-shot object detection) — замена ненадёжной CLIP-проверке категории `weapons`;
   - Комбинированный вердикт на основе нескольких моделей.
 - **Готовый REST API** с автоматической документацией **Swagger UI** (`/docs`) и **ReDoc** (`/redoc`).
 - **Автоматическая загрузка ML-моделей** при старте приложения (кэшируются локально).
@@ -57,6 +58,7 @@
 | Работа с переменными окружения | python-dotenv |
 | Нейросетевые модели | Transformers (Hugging Face), PyTorch |
 | Мультимодальная модель | OpenAI CLIP (ViT-B/32) |
+| Детектор оружия | OWLv2 (`google/owlv2-base-patch16-ensemble`) |
 | Морфология русского языка | pymorphy3 |
 | Обработка изображений | Pillow (PIL) |
 | Язык | Python 3.12 |
@@ -100,7 +102,9 @@ moderation-post/
     │   │       ├── image_classifier.py           # Классификатор NSFW изображений
     │   │       ├── image_prediction_result.py    # Результат классификации изображений
     │   │       ├── clip_classifier.py            # CLIP-классификатор по категориям
-    │   │       └── clip_prediction_result.py     # Результат CLIP-классификации
+    │   │       ├── clip_prediction_result.py     # Результат CLIP-классификации
+    │   │       ├── weapon_detector.py            # Детектор оружия (OWLv2)
+    │   │       └── weapon_prediction_result.py   # Результат детекции оружия
     │   ├── models/
     │   │   ├── dto/
     │   │   │   └── text_request.py      # DTO запроса модерации текста
@@ -183,6 +187,7 @@ DEBUG=true
 TEXT_MODEL=textdetox/bert-multilingual-toxicity-classifier
 IMAGE_MODEL=Falconsai/nsfw_image_detection
 IMAGE_CLIP_MODEL=openai/clip-vit-base-patch32
+WEAPON_MODEL=google/owlv2-base-patch16-ensemble
 MODEL_CACHE_DIR=./models
 DEVICE=auto
 HF_TOKEN=<ваш HF-токен>
@@ -230,6 +235,9 @@ fastapi run app/main.py --host 0.0.0.0 --port 8000
 | `NSFW_THRESHOLD` | float | `0.85` | Порог уверенности NSFW-модели |
 | `IMAGE_CLIP_MODEL` | str | `openai/clip-vit-base-patch32` | Мультимодальная модель CLIP |
 | `CLIP_THRESHOLD` | float | `0.75` | Порог уверенности CLIP-классификатора |
+| `WEAPON_MODEL` | str | `google/owlv2-base-patch16-ensemble` | Модель zero-shot детекции оружия (OWLv2) |
+| `WEAPON_THRESHOLD` | float | `0.40` | Порог уверенности детектора оружия |
+| `WEAPON_MIN_AREA_FRACTION` | float | `0.005` | Минимальная доля площади изображения для бокса (0.5%), отсекает мелкие ложные срабатывания |
 | `HF_TOKEN` | str | — | Токен Hugging Face (подписанный/приватные модели) |
 | `MODEL_CACHE_DIR` | str | `./models` | Каталог локального кэша моделей |
 | `DEVICE` | str | `auto` | Устройство инференса: `auto`, `cpu`, `cuda` и т.д. |
@@ -447,10 +455,10 @@ class ClipPredictionResult:
 - **`image_policy.py`** — множество `FORBIDDEN_IMAGE_LABELS` с запрещёнными категориями CLIP:
   - `a photo containing marijuana`
   - `a photo containing drugs`
-  - `a photo containing weapons`
-  - `a pornographic photo`
   - `a photo containing violence`
 
+  Порнография и оружие в этот список не входят: первая обрабатывается NSFW-моделью (шаг 1), второе — отдельным детектором на базе OWLv2 (шаг 2).
+
 ### DI-фабрики — `app/core/dependencies.py`
 
 Использует `@lru_cache`, что превращает фабрики в синглтоны:
@@ -460,9 +468,10 @@ class ClipPredictionResult:
 | `get_text_classifier` | `TextClassifier` |
 | `get_clip_classifier` | `ClipClassifier` |
 | `get_image_classifier` | `ImageClassifier` |
+| `get_weapon_detector` | `WeaponDetector` (OWLv2) |
 | `get_profanity_detector` | `ProfanityDetector` (словарь + лемматизатор) |
 | `get_text_moderation_service` | `TextModerationService` |
-| `get_image_moderation_service` | `ImageModerationService` (классификатор + валидатор + CLIP) |
+| `get_image_moderation_service` | `ImageModerationService` (классификатор + валидатор + CLIP + детектор оружия) |
 | `get_image_validator` | `ImageValidator` |
 
 > Файл `app/container.py` содержит аналогичный, но неиспользуемый «ручной» контейнер, создающий детектор нецензурной лексики. Основной путь внедрения зависимостей — через `core/dependencies.py`.
@@ -485,9 +494,15 @@ class ClipPredictionResult:
   - возвращает `ImagePredictionResult` с меткой из `id2label` (обычно `nsfw`/`normal`);
 - **`image/clip_classifier.py`** — `ClipClassifier`:
   - загружает CLIP-модель и процессор сразу в конструкторе (при первом обращении);
-  - задаёт список из 6 текстовых меток (`a normal photo` + 5 запрещённых);
+  - задаёт список из 4 текстовых меток (`a normal photo` + 3 запрещённые: наркотики, насилие). Порнография отдельно обрабатывается NSFW-моделью на шаге 1;
   - считает вероятности соответствия изображения каждой метке;
-  - собирает метки, превышающие `CLIP_THRESHOLD`, и возвращает `ClipPredictionResult`.
+  - собирает метки, превышающие `CLIP_THRESHOLD`, и возвращает `ClipPredictionResult`;
+- **`image/weapon_detector.py`** — `WeaponDetector`:
+  - загружает OWLv2-модель и процессор (zero-shot object detection) при первом обращении;
+  - ищет по 15 текстовым запросам: `gun`, `pistol`, `rifle`, `shotgun`, `knife`, `machete`, `sword`, `bomb`, `grenade` и др.;
+  - через `post_process_object_detection` получает боксы и уверенность по каждой запрошенной категории;
+  - отсеивает боксы меньше `WEAPON_MIN_AREA_FRACTION` от площади изображения и с уверенностью ниже `WEAPON_THRESHOLD`;
+  - возвращает `WeaponPredictionResult` с найденными категориями.
 
 ### Слой правил — `app/moderation/`
 
@@ -531,12 +546,15 @@ class ClipPredictionResult:
 1. **Валидация.** `ImageValidator.validate(image)` — проверка размеров (см. выше).
 2. **Шаг 1 — NSFW-модель.** `ImageClassifier.predict(image)`:
    - если метка `nsfw` и уверенность ≥ `NSFW_THRESHOLD` (0.85) → решение **отклонить**: `reason="NSFW"`, `detected_labels=["nsfw"]`;
-3. **Шаг 2 — CLIP-модель.** `ClipClassifier.predict(image)`:
-   - из меток, превысивших `CLIP_THRESHOLD`, отбираются только те, что входят в `FORBIDDEN_IMAGE_LABELS`;
+3. **Шаг 2 — детектор оружия.** `WeaponDetector.predict(image)`:
+   - zero-shot детекция по запросам оружия (OWLv2, боксы + уверенность);
+   - если найдена категория с уверенностью ≥ `WEAPON_THRESHOLD` и боксом больше `WEAPON_MIN_AREA_FRACTION` → решение **отклонить**: `reason="FORBIDDEN_CONTENT"`, `detected_labels=[...]`;
+4. **Шаг 3 — CLIP-модель.** `ClipClassifier.predict(image)`:
+   - из меток, превысивших `CLIP_THRESHOLD`, отбираются только те, что входят в `FORBIDDEN_IMAGE_LABELS` (наркотики, насилие);
    - если такие категории найдены → решение **отклонить**: `reason="FORBIDDEN_CONTENT"`, `detected_labels=[...]`;
-4. **Шаг 3 — «нормальное» изображение.** Если ни одно из правил не сработало → `approved=True`, `reason="OK"`, `label="normal"`, `score=1.0`.
+5. **Шаг 4 — «нормальное» изображение.** Если ни одно из правил не сработало → `approved=True`, `reason="OK"`, `label="normal"`, `score=1.0`.
 
-Таким образом, изображение одобряется только если оно прошло и NSFW-модель, и CLIP-проверку по запрещённым категориям.
+Таким образом, изображение одобряется только если оно прошло и NSFW-модель, и детектор оружия, и CLIP-проверку по запрещённым категориям.
 
 ---