70 lines
1.2 KiB
Python
70 lines
1.2 KiB
Python
import torch
|
|
from PIL import Image
|
|
|
|
from app.ml.image.image_prediction_result import ImagePredictionResult
|
|
from app.ml.model_manager import model_manager
|
|
|
|
|
|
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
|
|
)
|