added image ai-moderation

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

@@ -0,0 +1,83 @@
+from PIL import Image
+import torch
+
+from app.ml.model_manager import model_manager
+from app.ml.image.image_prediction_result import ImagePredictionResult
+
+
+class ImageClassifier:
+
+
+  def predict(
+      self,
+      image: Image.Image
+  ) -> ImagePredictionResult:
+
+
+    inputs = (
+      model_manager.image_processor(
+          image,
+          return_tensors="pt"
+      )
+    )
+
+
+    inputs = {
+      key: value.to(model_manager.device)
+      for key, value in inputs.items()
+    }
+
+
+    with torch.no_grad():
+
+      outputs = (
+        model_manager.image_model(**inputs)
+      )
+
+
+    probabilities = torch.softmax(
+        outputs.logits,
+        dim=1
+    )
+
+
+    raw_scores = probabilities[0].tolist()
+
+
+    predicted_index = (
+      torch.argmax(probabilities, dim=1)
+      .item()
+    )
+
+
+    labels = (
+      model_manager.image_model
+      .config
+      .id2label
+    )
+
+
+    label = labels[predicted_index]
+
+
+    score = raw_scores[predicted_index]
+
+
+    print("====================")
+    print("IMAGE RESULT")
+    print("Labels:", labels)
+    print("Label:", label)
+    print("Score:", score)
+    print("Scores:", raw_scores)
+
+
+    return ImagePredictionResult(
+
+        label=label,
+
+        score=score,
+
+        approved=False,
+
+        raw_scores=raw_scores
+    )