Spaces:
Running
Running
File size: 1,748 Bytes
bfb2e8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
import numpy as np
from model_loader import models
class Inferencer:
def __init__(self):
self.clip_model = models.clip_model
self.svm_model = models.svm_model
@torch.no_grad()
def predict(self, image_tensor:torch.Tensor) -> dict:
"""
Takes a preprocessed image tensor and returns the classification result.
Args:
image_tensor (torch.Tensor): The preprocessed image tensor.
Returns:
dict: A dictionary containing the classification label and confidence score.
"""
image_features = self.clip_model.encode_image(image_tensor)
image_features_np = image_features.cpu().numpy()
prediction = self.svm_model.predict(image_features_np)[0]
if hasattr(self.svm_model, "predict_proba"):
# If yes, use predict_proba for a true confidence score
confidence_scores = self.svm_model.predict_proba(image_features_np)[0]
confidence = float(np.max(confidence_scores))
else:
# If no, use decision_function as a fallback confidence measure.
# The absolute value of the decision function score indicates confidence.
# We can apply a sigmoid function to scale it to a [0, 1] range for consistency.
decision_score = self.svm_model.decision_function(image_features_np)[0]
confidence = 1 / (1 + np.exp(-np.abs(decision_score)))
confidence = float(confidence)
label_map = {0: 'real', 1: 'fake'}
classification_label = label_map.get(prediction, "unknown")
return {
"classification": classification_label,
"confidence": confidence
}
inferencer = Inferencer() |