Merge pull request #5 from SlimusMinus/change-image_model

change image model

ai-moderation/app/config/image_policy.py

@@ -3,8 +3,6 @@ FORBIDDEN_IMAGE_LABELS = {
   s.lower() for s in {
     "a photo containing marijuana",
     "a photo containing drugs",
-    "a photo containing weapons",
-    "a pornographic photo",
     "a photo containing violence",
   }
 }

ai-moderation/app/config/settings.py

@@ -21,6 +21,9 @@ class Settings(BaseSettings):
   NSFW_THRESHOLD: float = 0.85
   IMAGE_CLIP_MODEL: str = "openai/clip-vit-base-patch32"
   CLIP_THRESHOLD: float = 0.75
+  WEAPON_MODEL: str = "google/owlv2-base-patch16-ensemble"
+  WEAPON_THRESHOLD: float = 0.40
+  WEAPON_MIN_AREA_FRACTION: float = 0.005
   MODEL_CACHE_DIR: str = "./models"
 
   DEVICE: str = "auto"

ai-moderation/app/core/dependencies.py

@@ -2,6 +2,7 @@ from functools import lru_cache
 
 from app.ml.image.clip_classifier import ClipClassifier
 from app.ml.image.image_classifier import ImageClassifier
+from app.ml.image.weapon_detector import WeaponDetector
 from app.ml.text_classifier import TextClassifier
 from app.moderation.image.validator import ImageValidator
 from app.moderation.profanity.detector import ProfanityDetector
@@ -49,10 +50,16 @@ def get_image_classifier() -> ImageClassifier:
   return ImageClassifier()
 
 
+@lru_cache
+def get_weapon_detector() -> WeaponDetector:
+  return WeaponDetector()
+
+
 @lru_cache
 def get_image_moderation_service() -> ImageModerationService:
   return ImageModerationService(
       classifier=get_image_classifier(),
       validator=get_image_validator(),
-      clip_classifier=get_clip_classifier()
+      clip_classifier=get_clip_classifier(),
+      weapon_detector=get_weapon_detector()
   )

ai-moderation/app/ml/image/clip_classifier.py

@@ -11,8 +11,6 @@ class ClipClassifier:
     "a normal photo",
     "a photo containing marijuana",
     "a photo containing drugs",
-    "a photo containing weapons",
-    "a pornographic photo",
     "a photo containing violence"
   ]
 

ai-moderation/app/services/image_moderation_service.py

@@ -5,6 +5,7 @@ from app.config.settings import settings
 from app.ml.image.clip_classifier import ClipClassifier
 from app.ml.image.image_classifier import ImageClassifier
 from app.ml.image.image_prediction_result import ImagePredictionResult
+from app.ml.image.weapon_detector import WeaponDetector
 from app.moderation.image.validator import ImageValidator
 
 class ImageModerationService:
@@ -13,11 +14,13 @@ class ImageModerationService:
       self,
       classifier: ImageClassifier,
       validator: ImageValidator,
-      clip_classifier: ClipClassifier
+      clip_classifier: ClipClassifier,
+      weapon_detector: WeaponDetector
   ):
 
     self.classifier = classifier
     self.clip_classifier = clip_classifier
+    self.weapon_detector = weapon_detector
     self.validator = validator
 
   def moderate(
@@ -48,7 +51,30 @@ class ImageModerationService:
       return nsfw_prediction
 
     # =========================
-    # 2. CLIP MODEL
+    # 2. WEAPON DETECTOR
+    # =========================
+
+    weapon_prediction = (
+      self.weapon_detector.predict(image)
+    )
+
+    if weapon_prediction.detected_labels:
+      return ImagePredictionResult(
+
+          label=weapon_prediction.label,
+
+          score=weapon_prediction.score,
+
+          approved=False,
+
+          reason="FORBIDDEN_CONTENT",
+
+          detected_labels=weapon_prediction.detected_labels
+
+      )
+
+    # =========================
+    # 3. CLIP MODEL
     # =========================
 
     clip_prediction = (
@@ -82,7 +108,7 @@ class ImageModerationService:
       )
 
     # =========================
-    # 3. NORMAL IMAGE
+    # 4. NORMAL IMAGE
     # =========================
 
     return ImagePredictionResult(

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-проверку по запрещённым категориям.
 
 ---