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
    )