fixed code

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

@@ -1,17 +1,12 @@
 import torch
-
 from PIL import Image
-from transformers import (
-  CLIPProcessor,
-  CLIPModel
-)
+from transformers import CLIPProcessor, CLIPModel
 
-from app.ml.image.clip_prediction_result import ClipPredictionResult
 from app.config.settings import settings
+from app.ml.image.clip_prediction_result import ClipPredictionResult
 
 
 class ClipClassifier:
-
   LABELS = [
     "a normal photo",
     "a photo containing marijuana",
@@ -21,106 +16,46 @@ class ClipClassifier:
     "a photo containing violence"
   ]
 
-
   def __init__(self):
+    self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Loading CLIP on {self.device}")
 
-    self.device = (
-      "cuda"
-      if torch.cuda.is_available()
-      else "cpu"
-    )
-
-    print(
-        f"Loading CLIP on {self.device}"
-    )
-
-    self.model = CLIPModel.from_pretrained(
-        settings.IMAGE_CLIP_MODEL
-    )
-
-    self.processor = CLIPProcessor.from_pretrained(
-        settings.IMAGE_CLIP_MODEL
-    )
+    self.model = CLIPModel.from_pretrained(settings.IMAGE_CLIP_MODEL)
+    self.processor = CLIPProcessor.from_pretrained(settings.IMAGE_CLIP_MODEL)
 
     self.model.to(self.device)
-
     self.model.eval()
 
-
-  def predict(
-      self,
-      image: Image.Image
-  ) -> ClipPredictionResult:
-
-
+  def predict(self, image: Image.Image) -> ClipPredictionResult:
     inputs = self.processor(
         text=self.LABELS,
         images=image,
         return_tensors="pt",
         padding=True
     )
-
-
-    inputs = {
-      k: v.to(self.device)
-      for k, v in inputs.items()
-    }
-
+    inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
     with torch.no_grad():
-
       outputs = self.model(**inputs)
-
-      logits = (
-        outputs
-        .logits_per_image
-      )
-
-      probs = (
-        logits.softmax(dim=1)[0]
-      )
-
+      logits = outputs.logits_per_image
+      probs = logits.softmax(dim=1)[0]
 
     scores = {}
-
     detected_labels = []
 
-
-    THRESHOLD = 0.75
-
-
     for index, probability in enumerate(probs):
-
       label = self.LABELS[index]
-
       score = float(probability)
-
-
       scores[label] = score
 
-
-      if (
-          label != "a normal photo"
-          and score >= THRESHOLD
-      ):
-
+      if label != "a normal photo" and score >= settings.CLIP_THRESHOLD:
         detected_labels.append(label)
 
-
-
     max_index = torch.argmax(probs).item()
 
-
     return ClipPredictionResult(
-
         label=self.LABELS[max_index],
-
-        score=float(
-            probs[max_index]
-        ),
-
+        score=float(probs[max_index]),
         detected_labels=detected_labels,
-
         scores=scores
-
     )