Spaces:
Sleeping
Sleeping
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 | |
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() |