diff --git a/ai-moderation/app/config/image_policy.py b/ai-moderation/app/config/image_policy.py index aa0f1ba..f2273ab 100644 --- a/ai-moderation/app/config/image_policy.py +++ b/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", } } \ No newline at end of file diff --git a/ai-moderation/app/config/settings.py b/ai-moderation/app/config/settings.py index 977f504..2ab77b1 100644 --- a/ai-moderation/app/config/settings.py +++ b/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" diff --git a/ai-moderation/app/core/dependencies.py b/ai-moderation/app/core/dependencies.py index 14634cf..20d34eb 100644 --- a/ai-moderation/app/core/dependencies.py +++ b/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() ) \ No newline at end of file diff --git a/ai-moderation/app/ml/image/clip_classifier.py b/ai-moderation/app/ml/image/clip_classifier.py index 5098226..58a77b4 100644 --- a/ai-moderation/app/ml/image/clip_classifier.py +++ b/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" ] diff --git a/ai-moderation/app/services/image_moderation_service.py b/ai-moderation/app/services/image_moderation_service.py index dba528f..b44abf6 100644 --- a/ai-moderation/app/services/image_moderation_service.py +++ b/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( diff --git a/readme.md b/readme.md index 39ad86f..bb713ef 100644 --- a/readme.md +++ b/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-проверку по запрещённым категориям. ---