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(