change image model

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(